jax.extend.linear_util.transformation2#
- jax.extend.linear_util.transformation2 = functools.partial(<class 'functools.partial'>, <function transformation2>)[source]#
Adds one more transformation to a WrappedFun.
- Parameters:
gen – the transformation generator function
fun (WrappedFun) – a WrappedFun on which to apply the transformation
gen_static_args – static args for the generator function
- Return type: