| |
| |
|
|
|
|
|
|
|
|
|
|
|
|
| def transpose(model, blob_in, blob_out, use_cudnn=False, **kwargs): |
| """Transpose.""" |
| if use_cudnn: |
| kwargs['engine'] = 'CUDNN' |
| return model.net.Transpose(blob_in, blob_out, **kwargs) |
|
|
|
|
| def sum(model, blob_in, blob_out, **kwargs): |
| """Sum""" |
| return model.net.Sum(blob_in, blob_out, **kwargs) |
|
|
|
|
| def reduce_sum(model, blob_in, blob_out, **kwargs): |
| """ReduceSum""" |
| return model.net.ReduceSum(blob_in, blob_out, **kwargs) |
|
|
|
|
| def sub(model, blob_in, blob_out, **kwargs): |
| """Subtract""" |
| return model.net.Sub(blob_in, blob_out, **kwargs) |
|
|
|
|
| def mat_mul(model, blob_in, blob_out, **kwargs): |
| """Matrix multiplication""" |
| return model.net.MatMul(blob_in, blob_out, **kwargs) |
|
|
|
|
| def arg_min(model, blob_in, blob_out, **kwargs): |
| """ArgMin""" |
| return model.net.ArgMin(blob_in, blob_out, **kwargs) |
|
|
| def batch_mat_mul(model, blob_in, blob_out, |
| enable_tensor_core=False, **kwargs): |
| if enable_tensor_core: |
| kwargs['engine'] = 'TENSORCORE' |
|
|
| return model.net.BatchMatMul(blob_in, blob_out, **kwargs) |
|
|
| def sparse_lengths_sum_4bit_rowwise_sparse(model, blob_in, blob_out, **kwargs): |
| return model.net.SparseLengthsSum4BitRowwiseSparse(blob_in, blob_out, **kwargs) |
|
|
