MCPcopy Create free account
hub / github.com/NVIDIA/cutlass / _get_batch_stride

Method _get_batch_stride

python/cutlass_cppgen/op/gemm.py:546–559  ·  view source on GitHub ↗

Returns the batch stride of ``tensor``. If ``tensor`` is only rank-2, batch stride is 0. :param tensor: tensor object to process :type tensor: numpy/cupy/torch array/tensor object :return: stride between each matrix in the batch :rtype: int

(self, tensor)

Source from the content-addressed store, hash-verified

544 return max(A_batch, B_batch)
545
546 def _get_batch_stride(self, tensor) -> int:
547 """
548 Returns the batch stride of ``tensor``. If ``tensor`` is only rank-2, batch stride is 0.
549
550 :param tensor: tensor object to process
551 :type tensor: numpy/cupy/torch array/tensor object
552
553 :return: stride between each matrix in the batch
554 :rtype: int
555 """
556 if tensor is not None and len(tensor.shape) > 2:
557 return tensor.shape[-2] * tensor.shape[-1]
558 else:
559 return 0
560
561 def _get_problem_args(self, A, B, C, D) -> tuple:
562 """

Callers 1

runMethod · 0.95

Calls

no outgoing calls

Tested by

no test coverage detected