composing - piping operations

[jakubMitura14]

Hello, simple question can we safely pipe operations using, for example, toolz pipe.
This is just syntactic sugar that still makes code more readable

[jakevdp]

Yes, toolz.pipe should work correctly with JAX; for example:

import jax
import jax.numpy as jnp
import toolz

fun = jax.jit(
    lambda x: toolz.pipe(x,
      jnp.sin,
      lambda x: x + 1,
      jnp.exp))

x = jnp.arange(4)
print(fun(x))
# [2.7182817 6.305807  6.7483463 3.1302722]

If you hit any situations where it doesn't work, feel free to open an issue and we'll see if there's anything we can do on the JAX side to make it work!

[jakubMitura14]

Thank you @jakevdp ! Just to be sure also I seen and used multiple times partial with jax without a problem, but should I avoid it in some cases