Raj Thakur

icon indicating an email [email protected]

icon company @aws-neuron icon location Cupertino icon location Building Next Generation ML Accelerators for AWS

Results 4 comments of Raj Thakur

`all_to_all` operation generates invalid HLO

> I wonder whether this was caused by an OpenXLA pin update. Were you able to identify when this started happening? Looks unlikely as I was able to replicate with...

Remove excessive warn message in maybe_get_jax as it creates too many log lines during training

Since JAX is an optional dependency, we should introduce a new env variable that applies to maybe_get_jax so that the function return None if the env is not set. Additional...

Remove excessive warn message in maybe_get_jax as it creates too many log lines during training

@steveepreston The issue is resolved in PT/XLA 2.8.1 There were multiple reasons contributing to perf degradation, This being one., but the major issue was compile flag change specially in 2.8.0

PJRT_Client_Destroy is not invoked in PT/XLA 2.8.0+

Hi @zhanyong-wan @ysiraichi, As stated earlier, We've implemented proper cleanup mechanisms in our PJRT client that depend on the PJRT_Client_Destroy lifecycle call. The current implementation require explicit runtime shutdown to...