@@ -114,25 +114,6 @@ mutable struct TracedRNumber{T} <: RNumber{T}
114114 end
115115end
116116
117- struct XLAArray{T,N} <: RArray{T,N}
118- # size::NTuple{N,Int}
119- end
120-
121- mutable struct ConcreteRArray{T,N} <: RArray{T,N}
122- data:: XLA.AsyncBuffer
123- # data::XLAArray{T, N}
124- shape:: NTuple{N,Int}
125- end
126-
127- const WrappedConcreteRArray{T,N} = WrappedArray{T,N,ConcreteRArray,ConcreteRArray{T,N}}
128- const AnyConcreteRArray{T,N} = Union{ConcreteRArray{T,N},WrappedConcreteRArray{T,N}}
129-
130- mutable struct ConcreteRNumber{T} <: RNumber{T}
131- data:: XLA.AsyncBuffer
132- end
133-
134- const TracedType = Union{TracedRArray,TracedRNumber,MissingTracedValue}
135-
136117include (" Ops.jl" )
137118include (" TracedUtils.jl" )
138119
@@ -143,6 +124,8 @@ include("ConcreteRArray.jl")
143124
144125include (" linear_algebra.jl" )
145126
127+ const TracedType = Union{TracedRArray,TracedRNumber,MissingTracedValue}
128+
146129include (" ControlFlow.jl" )
147130include (" Tracing.jl" )
148131include (" Compiler.jl" )
@@ -163,7 +146,7 @@ function Enzyme.make_zero(
163146 return res
164147end
165148
166- using . Compiler: @compile , @code_hlo , @jit , create_result, compile
149+ using . Compiler: @compile , @code_hlo , @jit , traced_getfield, create_result, compile
167150export ConcreteRArray, ConcreteRNumber, @compile , @code_hlo , @jit , @trace
168151
169152const registry = Ref {MLIR.IR.DialectRegistry} ()
0 commit comments