ejkernel.xla_utils package

• ejkernel.xla_utils.cumsum
  • chunk_global_cumsum()
  • chunk_global_cumsum_scalar()
  • chunk_global_cumsum_vector()
  • chunk_local_cumsum()
  • chunk_local_cumsum_scalar()
  • chunk_local_cumsum_vector()
• ejkernel.xla_utils.shardings
  • get_corrected_named_sharding()
  • reorder_sequence()
• ejkernel.xla_utils.utils
  • cdiv()
  • identity_dtype_convert()
  • prepare_chunk_indices()
  • prepare_chunk_offsets()
  • prepare_cu_seqlens_from_mask()
  • prepare_lens()
  • prepare_lens_from_mask()
  • prepare_position_ids()
  • prepare_sequence_ids()
  • prepare_token_indices()