Jonatan Kłosko
Jonatan Kłosko
### Description ```python import jax import jax.numpy as jnp def f(x, idx): return jax.lax.dynamic_slice(x, [idx], [2]) x = jnp.array([1, 2, 3]) idx = jnp.array(2) # Print lowered HLO print(jax.jit(f).lower(x, idx).as_text())...
### Description ```python import jax import jax.numpy as jnp def f(x, update): return jax.lax.dynamic_update_slice(x, update, [1, 1]) x = jnp.array([[1, 2, 3], [4, 5, 6]]) update = jnp.array([[7, 8], [9,...
### Description ```python import jax import jax.numpy as jnp def f(x): return jax.lax.clz(x) x = jnp.array([0, 1 tensor "mhlo.return"(%3) : (tensor) -> () }) {base_dilations = dense : tensor, padding...
### Description ```python import jax import jax.numpy as jnp def f(x): return jax.lax.reduce_window( x, 0, jnp.add, window_dimensions=(2,), window_strides=(1,), padding=[(0, 0)], base_dilation=(1,), window_dilation=(1,)) x = jnp.array([1, 2, 4]) # Print lowered...
### Description ```python import jax import jax.numpy as jnp def f(x): def reducer(op_val_index, acc_val_index): op_val, op_index = op_val_index acc_val, acc_index = acc_val_index return (op_val + acc_val, op_index + acc_index) idx...
This integrates the PjRt plugin from the [`jax-metal`](https://developer.apple.com/metal/jax/) for running on the Apple GPU. To test it, one can set `client: :mps` on EXLA backend/compiler. Since the plugin is loaded...
This way tensors require less memory by default and people can opt-in if they work with really big numbers. We already default to 32-bit precision for floats. Sidenote: Jax does...
### Description ```python import jax import jax.numpy as jnp def f(x): def cond(carry): i, x, acc = carry return i < x.shape[0] def body(carry): i, x, acc = carry return...
We have the manually-triggered background job for computing ranks information (CAD). For each of the ranks tables, it removes all rows, then computes and inserts fresh rows. While the computation...