Skip to content

Add sketch of new Tensor interface#2457

Open
dennisYatunin wants to merge 16 commits into
mainfrom
dy/tensors
Open

Add sketch of new Tensor interface#2457
dennisYatunin wants to merge 16 commits into
mainfrom
dy/tensors

Conversation

@dennisYatunin
Copy link
Copy Markdown
Member

This PR refactors the Geometry module to use a much simpler interface, so that all vector/tensor operations can expressed using standard math operations instead of custom API functions. This will allow us to remove a large amount of duplicate code, reduce compilation latency, and speed up GPU runs by optimizing the geometry data passed to each kernel.

  • Code follows the style guidelines OR N/A.
  • Unit tests are included OR N/A.
  • Code is exercised in an integration test OR N/A.
  • Documentation has been added/updated OR N/A.

@dennisYatunin dennisYatunin force-pushed the dy/tensors branch 23 times, most recently from d1c85a9 to dc71352 Compare February 21, 2026 03:40
@dennisYatunin dennisYatunin force-pushed the dy/tensors branch 6 times, most recently from 271e3cd to 1f8a383 Compare February 25, 2026 06:57
Copy link
Copy Markdown
Member Author

@dennisYatunin dennisYatunin left a comment

Choose a reason for hiding this comment

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

I think this looks mergeable, save for some minor nitpicks. If you want to check this with ClimaAtmos before merging, feel free to define some aliases like const AxisTensor = Tensor, which I used to ensure backward compatibility when refactoring RecursiveApply.

Comment thread docs/src/matrix_fields.md Outdated
Comment thread ext/cuda/column_matrix_helpers.jl Outdated
Comment thread src/Geometry/conversions.jl
Comment thread src/Geometry/conversions.jl Outdated
Comment thread src/Geometry/conversions.jl Outdated
Comment thread src/Geometry/mul_with_projection.jl Outdated
Comment thread src/Geometry/tensors.jl Outdated
Comment on lines +658 to +660
@inline outer(x::AbstractVector, y::AbstractVector) = x * y'
@inline outer(x::AbstractVector, y::Number) = x * y
@inline outer(x::AbstractVector, y) = nested_broadcast(y -> x ⊗ y, y)
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Suggested change
@inline outer(x::AbstractVector, y::AbstractVector) = x * y'
@inline outer(x::AbstractVector, y::Number) = x * y
@inline outer(x::AbstractVector, y) = nested_broadcast(y -> x y, y)
@inline outer(x, y) = x * y'

Both * and ' already call nested_broadcast, and you can call ' on a Number (that's just a no-op).

Comment on lines +235 to +236
- nested type (Tuple or NamedTuple), scalar type (Number, SMatrix, or
Tensor{2}/adjoint thereof), nested type (Tuple or NamedTuple)
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

This also looks like a relic from before the RecursiveApply refactor.

Comment thread src/Topologies/dss_transform.jl Outdated
Comment on lines +48 to +52
@inline dss_transform(
arg::Geometry.Covariant3Vector,
local_geometry::Geometry.LocalGeometry,
weight,
) = arg * weight
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

This case is already covered by the method below.

Comment thread test/Geometry/method_info.jl Outdated
(Tensor{2,FT,Tuple{Covariant3Axis,WAxis},SMatrix{1,1,FT,1}},LG_123_XYZ{FT}, 103),
]
end

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I feel like testing the padded FLOP counts here doesn't serve much of a purpose. You can leave it in if you want, but we only need tests like this for sparse tensors, which won't be introduced until a future PR.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

I'll remove them, they don't make sense with the uniform padded storage.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Geometry module refactor

2 participants