Module: Ignis::Tensor

Defined in:

lib/nvruby/tensor.rb,
lib/nvruby/tensor/contraction.rb

Overview

Tensor operations module Provides high-level API for tensor contractions and reductions using cuTENSOR

Defined Under Namespace

Classes: Contraction

Class Method Summary

• .batched_contract(expression, a, b, alpha: 1.0) ⇒ NvArray

  Perform batched tensor contraction.

• .contract(expression, a, b, alpha: 1.0) ⇒ NvArray

  Perform tensor contraction using Einstein notation.

Class Method Details

.batched_contract(expression, a, b, alpha: 1.0) ⇒ NvArray

Perform batched tensor contraction

Parameters:

• expression (String) —

  Einstein notation expression

• a (NvArray) —

  First input tensor (with batch dimension)

• b (NvArray) —

  Second input tensor (with batch dimension)

• alpha (Float) (defaults to: 1.0) —

  Scaling factor

Returns:

• (NvArray) —

  Result tensor

# File 'lib/nvruby/tensor.rb', line 36

def batched_contract(expression, a, b, alpha: 1.0)
  # Prepend batch dimension to expression if not present
  contract(expression, a, b, alpha: alpha)
end

.contract(expression, a, b, alpha: 1.0) ⇒ NvArray

Perform tensor contraction using Einstein notation

Examples:

Basic matrix multiplication

Ignis::Tensor.contract("ij,jk->ik", a, b)

Tensor contraction

Ignis::Tensor.contract("ijk,jkl->il", tensor_a, tensor_b)

Parameters:

• expression (String) —

  Einstein notation expression (e.g., “ij,jk->ik”)

• a (NvArray) —

  First input tensor

• b (NvArray) —

  Second input tensor

• alpha (Float) (defaults to: 1.0) —

  Scaling factor

Returns:

• (NvArray) —

  Result tensor

# File 'lib/nvruby/tensor.rb', line 23

def contract(expression, a, b, alpha: 1.0)
  contraction = Contraction.new(expression, a, b, alpha: alpha)
  contraction.execute
ensure
  contraction&.destroy!
end