Mapping over dynamic number of rows in XLA

[gurgeh]

I've been reading https://jax.readthedocs.io/en/latest/notebooks/XLA_in_Python.html and have some working code, but the tutorial only shows examples where the number of rows (/observations/samples) is constant. How do I compile a function where one dimension is dynamic?

I though it might be connected to "dynamic dimensions", but when I use for example:

X_shape = xc.Shape.array_shape(np.dtype(np.float32), (4,3), dynamic_dimensions=[True, False])
X = xops.Parameter(c, 0, X_shape)
...
host_X = np.array([[3.0+e, 4.0+e, 5.0+e] for e in [0.1, 0.2, 0.3, 0.5]], dtype=np.float32)

I get
Invalid argument: Executable expected shape f32[<=4,3]{1,0} for argument 0 but got incompatible shape f32[4,3]{1,0}
which seems odd to me since 4 <= 4..

I have tried with xops.Map over a computation and without, but I can't get it to work.Do I really have to compile a new function depending on number of rows in input? Is the "dynamic dimensions" thing a red herring?

[jekbradbury]

This is really a question about the XLA dynamic padder, not JAX; while the documentation on the feature is lacking, you might find something useful in the tests.