law batched_plan_geometry

bendtorch-core@0.1.0.0 · src/tensor/batched_matmul_LAWS.bend:406

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/batched_matmul.bend as Batched

406# Whatever the operands' dimensions, a plan that succeeds has PyTorch's
407# output dimensions: a 1-D left operand promoted to a row and a 1-D right
408# one to a column, both dropped again, and the batches broadcast
409# right-aligned, an extent of 1 taking the other one's. It reads the left
410# operand as its promoted extents padded with ones to rank 4, the right one
411# likewise (their inner extents are equal), and the output's batch padded
412# to two; these are the extents the build's layouts and
413# `batched_build_logical_cells` read with.
414law batched_plan_geometry:
415  for +left: List<&2, U32>
416  for +right: List<&2, U32>
417  torch_geometry(Shape.naturals(left), Shape.naturals(right), Batched.plan_dims(left, right))