From 7e603920700660e28a05c42764cae1922ed2e40b Mon Sep 17 00:00:00 2001
From: Russell Power <power@google.com>
Date: Thu, 8 Jun 2023 11:27:21 -0700
Subject: [PATCH] Fix a few shape inference bugs.

PiperOrigin-RevId: 538843234
---
 tf_quant_finance/models/generic_ito_process_test.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tf_quant_finance/models/generic_ito_process_test.py b/tf_quant_finance/models/generic_ito_process_test.py
index db058632a..82ba92ce1 100644
--- a/tf_quant_finance/models/generic_ito_process_test.py
+++ b/tf_quant_finance/models/generic_ito_process_test.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Tests for `sample_paths` of `ItoProcess`."""
 
+import os
 from unittest import mock  # pylint: disable=g-importing-member
 from absl.testing import parameterized
 import numpy as np
@@ -748,4 +749,5 @@ def _gaussian(xs, variance):
 
 
 if __name__ == '__main__':
+  os.environ['TF_XLA_FLAGS'] = '--tf_mlir_enable_mlir_bridge=true'
   tf.test.main()