Add sketch of new Tensor interface by dennisYatunin · Pull Request #2457 · CliMA/ClimaCore.jl

dennisYatunin · 2026-02-14T06:28:56Z

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

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.

dennisYatunin · 2026-05-12T20:28:40Z


 - A `UniformScaling`, which contains a `Number`
- A `DiagonalMatrixRow`, which can contain either a `Number` or a tensor (represented as a `Geometry.Axis2Tensor`)
+- A `DiagonalMatrixRow`, which can contain either a `Number` or a rank-2 tensor (represented as a `Geometry.Tensor{2}`)


Should this say AbstractTensor{2} instead of Tensor{2}? Same for the instance of Tensor{2} below.

dennisYatunin · 2026-05-12T20:29:12Z

    x::T1,
    y::T2,
-) where {T1 <: Geometry.AdjointAxisVector, T2 <: Geometry.Axis2Tensor} = (x * y)'
+) where {T1 <: Geometry.AbstractCovector, T2 <: Geometry.Tensor{2}} = (x * y)'


Suggested change

) where {T1 <: Geometry.AbstractCovector, T2 <: Geometry.Tensor{2}} = (x * y)'

) where {T1 <: Geometry.AbstractCovector, T2 <: Geometry.AbstractTensor{2}} = (x * y)'

dennisYatunin · 2026-05-12T20:41:11Z

+# every conversion is a single matvec — names outside `lg`'s geometry `I`
+# ride the identity block of the padded matrix automatically. Same-type


This is a bit unclear. The LocalGeometry tensors now cover all available dimension names, and there's no additional padding here.

dennisYatunin · 2026-05-12T20:45:58Z

+# 1D vector types can be constructed from a scalar + LocalGeometry.
+# The LocalGeometry is ignored — the scalar is wrapped directly.


Does this need to be supported? It looks redundant.

dennisYatunin · 2026-05-12T20:48:28Z

-            lg.∂ξ∂x[2, 1]*v[1, 1]+lg.∂ξ∂x[2, 2]*v[2, 1] lg.∂ξ∂x[2, 1]*v[1, 2]+lg.∂ξ∂x[2, 2]*v[2, 2]
-        ]
-    )
+for (BT, VecType, fn) in (


It looks like BT should be a symbol here, and fn can be replaced with VecType.

dennisYatunin · 2026-05-12T22:13:06Z

+basis1(::Type{<:Tensor{2, <:Any, <:Tuple{B, Any}}}) where {B} = B
+basis2(::Type{<:Tensor{2, <:Any, <:Tuple{Any, B}}}) where {B} = B


Suggested change

basis1(::Type{<:Tensor{2, <:Any, <:Tuple{B, Any}}}) where {B} = B

basis2(::Type{<:Tensor{2, <:Any, <:Tuple{Any, B}}}) where {B} = B

basis1(::Type{T}) where {T <: AbstractTensor{2}} = Geometry.tensor_bases(T)[1]

basis1(::Type{T}) where {T <: AbstractTensor{2}} = Geometry.tensor_bases(T)[2]

Some of the other Tensor types below could also be replaced with AbstractTensors.

dennisYatunin · 2026-05-12T22:21:26Z

+@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)


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).

dennisYatunin · 2026-05-12T22:26:36Z

+- nested type (Tuple or NamedTuple), scalar type (Number, SMatrix, or
+  Tensor{2}/adjoint thereof), nested type (Tuple or NamedTuple)


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

dennisYatunin · 2026-05-12T22:30:56Z

+@inline dss_transform(
+    arg::Geometry.Covariant3Vector,
+    local_geometry::Geometry.LocalGeometry,
+    weight,
+) = arg * weight


This case is already covered by the method below.

dennisYatunin · 2026-05-12T22:39:52Z

+    (Tensor{2,FT,Tuple{Covariant3Axis,WAxis},SMatrix{1,1,FT,1}},LG_123_XYZ{FT}, 103),
    ]
 end



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.

dennisYatunin requested a review from nefrathenrici February 14, 2026 06:28

dennisYatunin force-pushed the dy/tensors branch 23 times, most recently from d1c85a9 to dc71352 Compare February 21, 2026 03:40

dennisYatunin force-pushed the dy/tensors branch 6 times, most recently from 271e3cd to 1f8a383 Compare February 25, 2026 06:57

nefrathenrici force-pushed the dy/tensors branch from 966abfc to 772a7fa Compare April 30, 2026 00:32

nefrathenrici marked this pull request as ready for review April 30, 2026 00:32

nefrathenrici force-pushed the dy/tensors branch 2 times, most recently from b6c3879 to b9f7f37 Compare April 30, 2026 19:40

nefrathenrici mentioned this pull request May 1, 2026

Update for ClimaCore geometry refactor CliMA/ClimaAtmos.jl#4455

Draft

nefrathenrici force-pushed the dy/tensors branch 7 times, most recently from 9c066c5 to e354450 Compare May 1, 2026 23:58

dennisYatunin and others added 5 commits May 6, 2026 13:03

Add sketch of new Tensor interface

13e1436

Refactor Geometry module

bcf8c9a

test [perf]

0666b57

fixes after rebase

9e1c1f4

fixes

522239a

nefrathenrici force-pushed the dy/tensors branch from e354450 to 062b580 Compare May 7, 2026 23:48

nefrathenrici linked an issue May 8, 2026 that may be closed by this pull request

Geometry module refactor #2399

Open

nefrathenrici mentioned this pull request May 8, 2026

Add true SCM atmos-only experiment in ClimaAtmos CliMA/ClimaAtmos.jl#3520

Open

5 tasks

zero pad [perf]

fbf6388

nefrathenrici force-pushed the dy/tensors branch from 062b580 to fbf6388 Compare May 8, 2026 21:01

nefrathenrici added 7 commits May 8, 2026 14:58

fixes

c3b7517

remove cob tensor

2c932ae

test main [perf]

a60c3e6

fies

d8cd436

last fixes

3c26dca

remove generated functions

4297e6d

update jet failures

61b3ac5

dennisYatunin commented May 12, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add sketch of new Tensor interface#2457

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

dennisYatunin commented Feb 14, 2026

Uh oh!

dennisYatunin left a comment

Uh oh!

dennisYatunin May 12, 2026

Uh oh!

dennisYatunin May 12, 2026

Uh oh!

dennisYatunin May 12, 2026

Uh oh!

dennisYatunin May 12, 2026

Uh oh!

dennisYatunin May 12, 2026

Uh oh!

dennisYatunin May 12, 2026

Uh oh!

dennisYatunin May 12, 2026

Uh oh!

dennisYatunin May 12, 2026

Uh oh!

dennisYatunin May 12, 2026

Uh oh!

dennisYatunin May 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	) where {T1 <: Geometry.AbstractCovector, T2 <: Geometry.Tensor{2}} = (x * y)'
	) where {T1 <: Geometry.AbstractCovector, T2 <: Geometry.AbstractTensor{2}} = (x * y)'

		# every conversion is a single matvec — names outside `lg`'s geometry `I`
		# ride the identity block of the padded matrix automatically. Same-type

		# 1D vector types can be constructed from a scalar + LocalGeometry.
		# The LocalGeometry is ignored — the scalar is wrapped directly.

		basis1(::Type{<:Tensor{2, <:Any, <:Tuple{B, Any}}}) where {B} = B
		basis2(::Type{<:Tensor{2, <:Any, <:Tuple{Any, B}}}) where {B} = B

		- nested type (Tuple or NamedTuple), scalar type (Number, SMatrix, or
		Tensor{2}/adjoint thereof), nested type (Tuple or NamedTuple)

Conversation

dennisYatunin commented Feb 14, 2026

Uh oh!

dennisYatunin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants