google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Functions for calculating masks of new cache items."""
17
import abc
18
from typing import Dict, Tuple
19

20
import tensorflow.compat.v2 as tf
21

22
from negative_cache import negative_cache
23

24

25
class CacheFilterFn(object, metaclass=abc.ABCMeta):
26

27
  @abc.abstractmethod
28
  def __call__(self, cache,
29
               new_items):
30
    pass
31

32

33
class IsInCacheFilterFn(CacheFilterFn):
34
  """Creates a mask for items that are already in the cache.
35

36
  Given a tuple of keys, this class is a function that checks if there is a
37
  cache element that matches exactly on all keys.
38
  """
39

40
  def __init__(self, keys):
41
    self.keys = keys
42

43
  def __call__(self, cache,
44
               new_items):
45
    datawise_matches = []
46
    for key in self.keys:
47
      cache_vals = cache.data[key]
48
      new_items_vals = new_items[key]
49
      if cache_vals.dtype.is_floating:
50
        raise NotImplementedError('Floating datatypes are not yet implemented.')
51
      cache_vals = tf.expand_dims(cache_vals, axis=0)
52
      new_items_vals = tf.expand_dims(new_items_vals, axis=1)
53
      elementwise = cache_vals == new_items_vals
54
      datawise = tf.reduce_all(elementwise, axis=range(2, tf.rank(elementwise)))
55
      datawise_matches.append(datawise)
56
    all_keys_datawise = tf.stack(datawise_matches, axis=2)
57
    all_keys_match = tf.reduce_all(all_keys_datawise, axis=2)
58
    in_cache = tf.reduce_any(all_keys_match, axis=1)
59
    return tf.logical_not(in_cache)
60