ejkernel API Reference

• ejkernel.callib package
  • ejkernel.callib._ejit
  • ejkernel.callib._pallas_call
  • ejkernel.callib._triton_call
  • ejkernel.callib._utils
• ejkernel.kernels package
  • ejkernel.kernels._pallas package
  • ejkernel.kernels._triton package
  • ejkernel.kernels._xla package
  • ejkernel.kernels._registry
• ejkernel.modules package
  • ejkernel.modules.operations package
  • ejkernel.modules.base
• ejkernel.ops package
  • ejkernel.ops.config package
  • ejkernel.ops.core package
  • ejkernel.ops.execution package
  • ejkernel.ops.utils package
  • ejkernel.ops.registry
• ejkernel.types package
  • ejkernel.types.mask
• ejkernel.xla_utils package
  • ejkernel.xla_utils.cumsum
  • ejkernel.xla_utils.shardings
  • ejkernel.xla_utils.utils
• ejkernel.benchmarks
• ejkernel.loggings
  • ColorFormatter
  • LazyLogger
  • ProgressLogger
  • create_step_profiler()
  • extinguish_profiler()
  • get_logger()
  • ignite_profiler()
• ejkernel.utils
  • arch_supports_fp8()
  • assert_close()
  • barrier_sync()
  • calculate_blocksize_and_wraps()
  • cdiv()
  • dtype_index()
  • generate_block_indices()
  • get_abs_err()
  • get_err_ratio()
  • get_gpu_arch()
  • get_input_shapes()
  • get_padded_headsize()
  • get_qkv_shardings()
  • get_segments_shardings()
  • get_sharding()
  • get_stride()
  • get_strides()
  • get_tpu_generation()
  • is_cdna()
  • is_fp8()
  • is_hip()
  • is_rdna()
  • kw_strides()
  • make_dummy_rpa_inputs()
  • make_mesh()
  • narrow()
  • next_power_of_2()
  • numeric_gen()
  • random_dense()
  • safe_autotune()
  • strides_from_shape()