easydel.kernels.matmul

easydel.kernels.matmul.custom_dot_general_kernel(lhs: Array, rhs: Array, dimension_numbers: Optional[Tuple[Tuple[Sequence[int], Sequence[int]], Tuple[Sequence[int], Sequence[int]]]] = None, precision=None, preferred_element_type=None, *args, **kwargs)

easydel.kernels.matmul.matmul(A: Array, B: Array, *, blocksize_m: Optional[int] = None, blocksize_k: Optional[int] = None, blocksize_n: Optional[int] = None, precision: Union[None, str, Precision, tuple[str, str], tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision], DotAlgorithm, DotAlgorithmPreset] = None, **_)

easydel.kernels.matmul.matmul_benchmark(unused_args=None)

easydel.kernels.matmul.matmul_grad_test()

easydel.kernels.matmul.matmul_test()

easydel.kernels.matmul.replace_dot_general_with_matmul()

easydel.kernels.matmul.test_dot_general_replacer()