GPU callbacks.

Commit

275 days ago

[jaxlib] Pack/unpack subbyte types to/from numpy arrays to support int2, uint2, int4, uint4, float4_e2m1fn subbyte types in CPU/GPU callbacks. Reverts a099b285307508efad12a015d6f6d9d13ae49077 PiperOrigin-RevId: 753974925

References

#27806 - [jaxlib] Pack/unpack subbyte types to/from numpy arrays to support int2, uint2, int4, uint4, float4_e2m1fn subbyte types in CPU/GPU callbacks.

Author

danielsuo

Committer

Google-ML-Automation

Parents

ff1672b3

jax e64049f4 - [jaxlib] Pack/unpack subbyte types to/from numpy arrays to support int2, uint2, int4, uint4, float4_e2m1fn subbyte types in CPU/GPU callbacks.

jax
e64049f4 - [jaxlib] Pack/unpack subbyte types to/from numpy arrays to support int2, uint2, int4, uint4, float4_e2m1fn subbyte types in CPU/GPU callbacks.