-
Notifications
You must be signed in to change notification settings - Fork 25
Open
Description
module @reactant_gesummv attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<f32> {enzymexla.memory_effects = []}, %arg1: tensor<f32> {enzymexla.memory_effects = []}, %arg2: tensor<64x64xf32> {enzymexla.memory_effects = []}, %arg3: tensor<64x64xf32> {enzymexla.memory_effects = []}, %arg4: tensor<64xf32> {enzymexla.memory_effects = []}) -> tensor<64xf32> attributes {enzymexla.memory_effects = []} {
%0 = stablehlo.dot_general %arg3, %arg4, contracting_dims = [0] x [0] : (tensor<64x64xf32>, tensor<64xf32>) -> tensor<64xf32>
%1 = stablehlo.dot_general %arg2, %arg4, contracting_dims = [0] x [0] : (tensor<64x64xf32>, tensor<64xf32>) -> tensor<64xf32>
%2 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<64xf32>
%3 = stablehlo.multiply %2, %1 : tensor<64xf32>
%4 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<64xf32>
%5 = stablehlo.multiply %4, %0 : tensor<64xf32>
%6 = stablehlo.add %3, %5 : tensor<64xf32>
return %6 : tensor<64xf32>
}
}this is computing
Metadata
Metadata
Assignees
Labels
No labels