[keras] layer: init Blackout layer - can be used to zero certain parts of tensors

keras.py

@@ -2,7 +2,7 @@
 import itertools
 from functools import cached_property
 import gc
-from collections import OrderedDict, defaultdict
+from collections import OrderedDict, defaultdict, Iterable
 import fnmatch
 import math
 import datetime
@@ -177,6 +177,26 @@ class KFeed(object):
         )
 
 
+def Blackout(ref, slic, name=None):
+    """Blackout (set to zero) values in tensor accoring to numpy like slices. Always done alogn zeroth axis.
+
+    Args:
+        ref (np.array): reference array for shape and dtype
+        slic (Iterable): iterable of numpy like slices
+        name (_type_, optional): _description_. Defaults to None.
+
+    Returns:
+        _type_: _description_
+    """
+    shape = ref.shape[1:]
+    dtype = ref.dtype
+    mask = np.ones(shape, dtype=dtype)
+    slic = slic if isinstance(slic, Iterable) else (slic,)
+    for _slic in slic:
+        mask[_slic] = 0
+    return tf.keras.layers.Lambda((lambda x: x * mask), name=name)
+
+
 def Normal(ref, const=None, ignore_zeros=False, name=None, **kwargs):
     """
     Normalizing layer according to ref.