Fix MultiVariateNormalDiag and similar distributions leaking tracers.

distrax/_src/bijectors/diag_linear.py

@@ -57,11 +57,21 @@ def __init__(self, diag: Array):
         batch_shape=diag.shape[:-1],
         dtype=diag.dtype)
     self._diag = diag
-    self.forward = self._bijector.forward
-    self.forward_log_det_jacobian = self._bijector.forward_log_det_jacobian
-    self.inverse = self._bijector.inverse
-    self.inverse_log_det_jacobian = self._bijector.inverse_log_det_jacobian
-    self.inverse_and_log_det = self._bijector.inverse_and_log_det
+
+  def forward(self, x: Array) -> Array:
+    return self._bijector.forward(x)
+
+  def forward_log_det_jacobian(self, x: Array) -> Array:
+    return self._bijector.forward_log_det_jacobian(x)
+
+  def inverse(self, y: Array) -> Array:
+    return self._bijector.inverse(y)
+
+  def inverse_log_det_jacobian(self, y: Array) -> Array:
+    return self._bijector.inverse_log_det_jacobian(y)
+
+  def inverse_and_log_det(self, y: Array) -> Tuple[Array, Array]:
+    return self._bijector.inverse_and_log_det(y)
 
   @property
   def diag(self) -> Array:

distrax/_src/bijectors/diag_plus_low_rank_linear.py

@@ -201,11 +201,21 @@ def __init__(self, diag: Array, u_matrix: Array, v_matrix: Array):
     self._diag = diag
     self._u_matrix = u_matrix
     self._v_matrix = v_matrix
-    self.forward = self._bijector.forward
-    self.forward_log_det_jacobian = self._bijector.forward_log_det_jacobian
-    self.inverse = self._bijector.inverse
-    self.inverse_log_det_jacobian = self._bijector.inverse_log_det_jacobian
-    self.inverse_and_log_det = self._bijector.inverse_and_log_det
+
+  def forward(self, x: Array) -> Array:
+    return self._bijector.forward(x)
+
+  def forward_log_det_jacobian(self, x: Array) -> Array:
+    return self._bijector.forward_log_det_jacobian(x)
+
+  def inverse(self, y: Array) -> Array:
+    return self._bijector.inverse(y)
+
+  def inverse_log_det_jacobian(self, y: Array) -> Array:
+    return self._bijector.inverse_log_det_jacobian(y)
+
+  def inverse_and_log_det(self, y: Array) -> Tuple[Array, Array]:
+    return self._bijector.inverse_and_log_det(y)
 
   @property
   def diag(self) -> Array: