Move op methods into dedicated objects and export them to torch namespace
sbrunk opened this issue · 5 comments
torch.scala is getting a bit large, and we still need to implement a bunch of ops. So I think it makes sense to move the op methods into smaller units and then define aliases via export clauses in the torch namespace/package like we already do with torch.special.
The difference being that we should do it also for methods that are not aliased explicitly in PyTorch (methods that are defined in the torch namespace there). A natural way to split them, is to use the ops structure in the PyTorch docs. We already group into these op groups in torch.scala already so it shouldn't be hard to move things out.
The advantage should be more clarity, maintainability and smaller compilation units.
Here's an idea for a structure:
torch/ops/PointwiseOps.scala
package torch.ops
object PointwiseOps:
def abs[D <: NumericNN](input: Tensor[D]): Tensor[D] = ...torch/ops/ReductionOps.scala
package torch.ops
object ReductionOps:
def argmax[D <: DType]( ...torch/torch.scala
export ops.PointwiseOps.*
export ops.ReductionOps.*I just did a quick feasibility test and there seems an issue with default methods with default parameters which are lost when an export is imported in a different compilation unit. See scala/scala3#17930. I could not find a way around it. Since we're we're using default parameters extensively, this is a blocker and I guess we'll have to wait until it is fixed.
Totally agree! that file is indeed getting pretty large.
While that bug is fixed I am thinking we could do try these workarounds (unsure if there are any negative implications of either)
-
Split them in files but still define them as part of
package torch- so noexportis needed -
Implementing traits and having a package object that extends all those traits (I've never done this myself but have seen something like this in other projects before)
Yesterday I spent a bit of time splitting them in files using option 1 - as well as the test suite. Everything seems to be working fine. I can upload a draft of my changes later today when I get home.
Btw, thanks a lot for openly sharing your thoughts and progress with the Reduction Ops in order to avoid any duplicate work. Will try to do the same.
Ah great I didn't think of 1. as a workaround. If you've already tried that successfully, let's go for it!
I tried 2. but ran into some weird circular issues when compiling the package object because the traits were also somewhere under torch. Also package objects, especially with inheritance, are (kind of) deprecated in Scala 3 as far as I know.
Btw, thanks a lot for openly sharing your thoughts and progress with the Reduction Ops in order to avoid any duplicate work. Will try to do the same.
Yeah I'm still trying to get into the habit of writing these things down more often. I think it's good exercise for me 😄
Let's keep this open even though we've split ops things into different files while defining them in the torch package in #30.
The export variant has a few advantages like being able to put ops into actual objects which in turn enables Scaladoc macros to work so I'd still like to try this option when it becomes possible.
You're right, for some reason 2. (traits and package objects inheriting from them) did work after all. Not sure what caused the issues last time I tried.