search

stanford-crfm / haliax

Named Tensors for Legible Deep Learning in JAX
Apache License 2.0
141 stars 9 forks source link

added implementation and tests for ravel function #33

Closed rohan-mehta-1024 closed 10 months ago

rohan-mehta-1024 commented 10 months ago

Description

Ravel function was added according to following syntax arr.ravel(new_axis_name: AxisSelector).

Fixes Issues

18

Unit test coverage

Minimal testing since this basically just calls hax.flatten_axes under the hood, but does check out when compared to jnp.ravel.

Known breaking changes/behaviors

N/A.