Skip to content

Change zipvmap to allow returning tuples of tensors #141

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
benikm91 wants to merge 1 commit into
base: main
Choose a base branch
from
+44 −7
Open

Change zipvmap to allow returning tuples of tensors #141

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 44 additions & 7 deletions core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,41 @@ import me.shadaj.scalapy.py.SeqConverters
import me.shadaj.scalapy.readwrite.Reader
import me.shadaj.scalapy.readwrite.Writer

trait ZipVmapResult[L: Label, FOut, MOut]:
type MappedOut
def toPy(out: FOut): py.Dynamic
def fromPy(pyOut: py.Dynamic): MOut

object ZipVmapResult:

// Single Tensor case
given singleTensor[L: Label, Shape <: Tuple: Labels, V]: ZipVmapResult[L, Tensor[Shape, V], Tensor[L *: Shape, V]] with
def toPy(out: Tensor[Shape, V]): py.Dynamic = out.jaxValue
def fromPy(pyOut: py.Dynamic): Tensor[L *: Shape, V] = Tensor(pyOut)

// Empty Tuple case
given emptyTuple[L: Label]: ZipVmapResult[L, EmptyTuple, EmptyTuple] with
def toPy(out: EmptyTuple): py.Dynamic = py.Dynamic.global.tuple(Seq.empty[py.Dynamic].toPythonProxy)
def fromPy(pyOut: py.Dynamic): EmptyTuple = EmptyTuple

// Inductive Tuple case (Pairs, Triples, N-tuples)
given consTuple[L: Label, H, HOut, T <: Tuple, TOut <: Tuple](using
hRes: ZipVmapResult[L, H, HOut],
tRes: ZipVmapResult[L, T, TOut]
): ZipVmapResult[L, H *: T, HOut *: TOut] with

def toPy(out: H *: T): py.Dynamic =
val headPy = hRes.toPy(out.head)
val tailSeq = tRes.toPy(out.tail).as[Seq[py.Dynamic]]
py.Dynamic.global.tuple((headPy +: tailSeq).toPythonProxy)

def fromPy(pyOut: py.Dynamic): HOut *: TOut =
val seq = pyOut.as[Seq[py.Dynamic]]
val h = hRes.fromPy(seq.head)
val tailPy = py.Dynamic.global.tuple(seq.tail.toPythonProxy)
val t = tRes.fromPy(tailPy)
h *: t

object FunctionalOps:

object ZipVmap:
Expand Down Expand Up @@ -51,33 +86,35 @@ object FunctionalOps:
* ...
* }
*/
def zipvmap[L: Label, Inputs <: Tuple, OutShape <: Tuple: Labels, OutV](
def zipvmap[L: Label, Inputs <: Tuple, FOut, MOut](

@marcelluethi marcelluethi Aug 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Type MOut is a free Parameter. It must be possible to deduce it directly from the Output type of the function (using match types)

@benikm91 benikm91 Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't we do this in ZipVmapResult? For each case, e.g., for single Tensor its:

ZipVmapResult[L, Tensor[Shape, V], Tensor[L *: Shape, V]]

@marcelluethi marcelluethi Aug 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe, but there should not be a type parameter the user sees that has no obvious function. L, Inputs and FOut is really all that is needed to determine the output Type. These are, at least for the trained user, relatively straight-forward to interpret. In contrast to that, MOut is a parameter that seems unrelated to anything the user inputs and hence hard to understand. Furthermore, I suspect that this makes it also harder for type inference to work correctly.

axis: Axis[L]
)(
tensors: Inputs // This is a Tuple of Tensors
tensors: Inputs
)(using
ev: SharedAxisRemover[ShapesOf[Inputs], L]
)(
f: TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]] => Tensor[OutShape, OutV]
): Tensor[L *: OutShape, OutV] =
f: TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]] => FOut
)(using
outMapper: ZipVmapResult[L, FOut, MOut]
): MOut =
val fpy = (args: py.Dynamic) =>
OnError.traceStack:
val tensorList = args.as[Seq[py.Dynamic]].zip(ev.shapesLabels).map: (jaxArr, labels) =>
Tensor(jaxArr)(using LabelsImpl(labels))

val inputTuple = Tuple.fromArray(tensorList.toArray)
val result = f(inputTuple.asInstanceOf[TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]]])
result.jaxValue
outMapper.toPy(result)

val jaxInputs = py.Dynamic.global.tuple(tensors.toArray.map(_.asInstanceOf[Tensor[?, ?]].jaxValue).toPythonProxy)
val indicesAsTuple = py.Dynamic.global.tuple(ev.indices.toPythonProxy)

val jaxResult = Jax.jax_helper.zipvmap(
fpy,
indicesAsTuple
)(jaxInputs)

Tensor(jaxResult)

outMapper.fromPy(jaxResult)
export ZipVmap.zipvmap

extension [T <: Tuple: Labels, V](t: Tensor[T, V])
Expand Down
Loading