From 18080bf07f54eac5b25b8fbf3655f20a910384f9 Mon Sep 17 00:00:00 2001
From: Edward2 Team <no-reply@google.com>
Date: Thu, 1 Apr 2021 02:00:05 -0700
Subject: [PATCH] Fix ECE metric update function syntax inside ensembles loop.
 Add MultiHeadMCSoftmaxDenseFA name in layers/__init__.py

PiperOrigin-RevId: 366208245
---
 edward2/tensorflow/layers/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/edward2/tensorflow/layers/__init__.py b/edward2/tensorflow/layers/__init__.py
index d9d8e9e5..2a374999 100644
--- a/edward2/tensorflow/layers/__init__.py
+++ b/edward2/tensorflow/layers/__init__.py
@@ -118,6 +118,7 @@
     "MCSigmoidDenseFA",
     "MCSoftmaxDense",
     "MCSoftmaxDenseFA",
+    "MultiHeadMCSoftmaxDenseFA",
     "MixtureLogistic",
     "NCPCategoricalPerturb",
     "NCPNormalOutput",