diffsl 0.11.4

A compiler for a domain-specific language for ordinary differential equations (ODE).
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

// RUN: %eopt --enzyme-batch %s | FileCheck %s

module {
  func.func @square(%x : tensor<3xf64>) -> tensor<3xf64> {
    %y = math.sin %x : tensor<3xf64>
    return %y : tensor<3xf64>
  }
  func.func @dsq(%x : tensor<10x2x3xf64>) -> tensor<10x2x3xf64> {
    %r = enzyme.batch @square(%x) { batch_shape=array<i64: 10, 2> } : (tensor<10x2x3xf64>) -> (tensor<10x2x3xf64>)
    return %r : tensor<10x2x3xf64>
  }
}

// CHECK:  func.func @dsq(%arg0: tensor<10x2x3xf64>) -> tensor<10x2x3xf64> {
// CHECK-NEXT:    %0 = call @batched_square(%arg0) : (tensor<10x2x3xf64>) -> tensor<10x2x3xf64>
// CHECK-NEXT:    return %0 : tensor<10x2x3xf64>
// CHECK-NEXT:  }
// CHECK:  func.func private @batched_square(%arg0: tensor<10x2x3xf64>) -> tensor<10x2x3xf64> {
// CHECK-NEXT:    %0 = math.sin %arg0 : tensor<10x2x3xf64>
// CHECK-NEXT:    return %0 : tensor<10x2x3xf64>
// CHECK-NEXT:  }