def get_params(img, scale, ratio):
        """Get parameters for ``crop`` for a random sized crop.

        Args:
            img (PIL Image): Image to be cropped.
            scale (tuple): range of size of the origin size cropped
            ratio (tuple): range of aspect ratio of the origin aspect ratio cropped

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
                sized crop.
        """
        area = img.size[0] * img.size[1]

        for attempt in range(10):
            target_area = random.uniform(*scale) * area
            log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
            aspect_ratio = math.exp(random.uniform(*log_ratio))

            w = int(round(math.sqrt(target_area * aspect_ratio)))
            h = int(round(math.sqrt(target_area / aspect_ratio)))

            if w <= img.size[0] and h <= img.size[1]:
                i = random.randint(0, img.size[1] - h)
                j = random.randint(0, img.size[0] - w)
                return i, j, h, w

        # Fallback to central crop
        in_ratio = img.size[0] / img.size[1]
        if (in_ratio < min(ratio)):
            w = img.size[0]
            h = w / min(ratio)
        elif (in_ratio > max(ratio)):
            h = img.size[1]
            w = h * max(ratio)
        else:  # whole image
            w = img.size[0]
            h = img.size[1]
        i = (img.size[1] - h) // 2
        j = (img.size[0] - w) // 2
        return i, j, h, w

def
get_params
(
img
,
scale
,
ratio
)
:
area
=
img
.
size
[
0
]
*
img
.
size
[
1
]
for
attempt
in
range
(
10
)
:
target_area
=
random
.
uniform
(
*
scale
)
*
area
log_ratio
=
(
math
.
log
(
ratio
[
0
]
)
,
math
.
log
(
ratio
[
1
]
)
)
aspect_ratio
=
math
.
exp
(
random
.
uniform
(
*
log_ratio
)
)
w
=
int
(
round
(
math
.
sqrt
(
target_area
*
aspect_ratio
)
)
)
h
=
int
(
round
(
math
.
sqrt
(
target_area
/
aspect_ratio
)
)
)
if
w
<=
img
.
size
[
0
]
and
h
<=
img
.
size
[
1
]
:
i
=
random
.
randint
(
0
,
img
.
size
[
1
]
h
)
j
=
random
.
randint
(
0
,
img
.
size
[
0
]
w
)
return
i
,
j
,
h
,
w
# Fallback to central crop
in_ratio
=
img
.
size
[
0
]
/
img
.
size
[
1
]
if
(
in_ratio
<
min
(
ratio
)
)
:
w
=
img
.
size
[
0
]
h
=
w
/
min
(
ratio
)
elif
(
in_ratio
>
max
(
ratio
)
)
:
h
=
img
.
size
[
1
]
w
=
h
*
max
(
ratio
)
else
:
# whole image
w
=
img
.
size
[
0
]
h
=
img
.
size
[
1
]
i
=
(
img
.
size
[
1
]
h
)
//
2
j
=
(
img
.
size
[
0
]
w
)
//
2
return
i
,
j
,
h
,
w

Get parameters for crop for a random sized crop .


def get_params(brightness, contrast, saturation, hue):
        """Get a randomized transform to be applied on image.

        Arguments are same as that of __init__.

        Returns:
            Transform which randomly adjusts brightness, contrast and
            saturation in a random order.
        """
        transforms = []

        if brightness is not None:
            brightness_factor = random.uniform(brightness[0], brightness[1])
            transforms.append(Lambda(lambda img: F.adjust_brightness(img, brightness_factor)))

        if contrast is not None:
            contrast_factor = random.uniform(contrast[0], contrast[1])
            transforms.append(Lambda(lambda img: F.adjust_contrast(img, contrast_factor)))

        if saturation is not None:
            saturation_factor = random.uniform(saturation[0], saturation[1])
            transforms.append(Lambda(lambda img: F.adjust_saturation(img, saturation_factor)))

        if hue is not None:
            hue_factor = random.uniform(hue[0], hue[1])
            transforms.append(Lambda(lambda img: F.adjust_hue(img, hue_factor)))

        random.shuffle(transforms)
        transform = Compose(transforms)

        return transform

def get_params ( brightness , contrast , saturation , hue ) : transforms = [ ] if brightness is not None : brightness_factor = random . uniform ( brightness [ 0 ] , brightness [ 1 ] ) transforms . append ( Lambda ( lambda img : F . adjust_brightness ( img , brightness_factor ) ) ) if contrast is not None : contrast_factor = random . uniform ( contrast [ 0 ] , contrast [ 1 ] ) transforms . append ( Lambda ( lambda img : F . adjust_contrast ( img , contrast_factor ) ) ) if saturation is not None : saturation_factor = random . uniform ( saturation [ 0 ] , saturation [ 1 ] ) transforms . append ( Lambda ( lambda img : F . adjust_saturation ( img , saturation_factor ) ) ) if hue is not None : hue_factor = random . uniform ( hue [ 0 ] , hue [ 1 ] ) transforms . append ( Lambda ( lambda img : F . adjust_hue ( img , hue_factor ) ) ) random . shuffle ( transforms ) transform = Compose ( transforms ) return transform

Get a randomized transform to be applied on image .


def get_params(degrees, translate, scale_ranges, shears, img_size):
        """Get parameters for affine transformation

        Returns:
            sequence: params to be passed to the affine transformation
        """
        angle = random.uniform(degrees[0], degrees[1])
        if translate is not None:
            max_dx = translate[0] * img_size[0]
            max_dy = translate[1] * img_size[1]
            translations = (np.round(random.uniform(-max_dx, max_dx)),
                            np.round(random.uniform(-max_dy, max_dy)))
        else:
            translations = (0, 0)

        if scale_ranges is not None:
            scale = random.uniform(scale_ranges[0], scale_ranges[1])
        else:
            scale = 1.0

        if shears is not None:
            shear = random.uniform(shears[0], shears[1])
        else:
            shear = 0.0

        return angle, translations, scale, shear

def get_params ( degrees , translate , scale_ranges , shears , img_size ) : angle = random . uniform ( degrees [ 0 ] , degrees [ 1 ] ) if translate is not None : max_dx = translate [ 0 ] * img_size [ 0 ] max_dy = translate [ 1 ] * img_size [ 1 ] translations = ( np . round ( random . uniform ( - max_dx , max_dx ) ) , np . round ( random . uniform ( - max_dy , max_dy ) ) ) else : translations = ( 0 , 0 ) if scale_ranges is not None : scale = random . uniform ( scale_ranges [ 0 ] , scale_ranges [ 1 ] ) else : scale = 1.0 if shears is not None : shear = random . uniform ( shears [ 0 ] , shears [ 1 ] ) else : shear = 0.0 return angle , translations , scale , shear


def download(self):
        """Download and extract the tarball, and download each individual photo."""
        import tarfile

        if self._check_integrity():
            print('Files already downloaded and verified')
            return

        download_url(self.url, self.root, self.filename, self.md5_checksum)

        # Extract file
        with tarfile.open(os.path.join(self.root, self.filename), 'r:gz') as tar:
            tar.extractall(path=self.root)

        # Download individual photos
        with open(os.path.join(self.root, 'dataset', 'SBU_captioned_photo_dataset_urls.txt')) as fh:
            for line in fh:
                url = line.rstrip()
                try:
                    download_url(url, os.path.join(self.root, 'dataset'))
                except OSError:
                    # The images point to public images on Flickr.
                    # Note: Images might be removed by users at anytime.
                    pass

Download and extract the tarball and download each individual photo .


def download(self):
        """Download the MNIST data if it doesn't exist in processed_folder already."""

        if self._check_exists():
            return

        makedir_exist_ok(self.raw_folder)
        makedir_exist_ok(self.processed_folder)

        # download files
        for url in self.urls:
            filename = url.rpartition('/')[2]
            file_path = os.path.join(self.raw_folder, filename)
            download_url(url, root=self.raw_folder, filename=filename, md5=None)
            self.extract_gzip(gzip_path=file_path, remove_finished=True)

        # process and save as torch files
        print('Processing...')

        training_set = (
            read_image_file(os.path.join(self.raw_folder, 'train-images-idx3-ubyte')),
            read_label_file(os.path.join(self.raw_folder, 'train-labels-idx1-ubyte'))
        )
        test_set = (
            read_image_file(os.path.join(self.raw_folder, 't10k-images-idx3-ubyte')),
            read_label_file(os.path.join(self.raw_folder, 't10k-labels-idx1-ubyte'))
        )
        with open(os.path.join(self.processed_folder, self.training_file), 'wb') as f:
            torch.save(training_set, f)
        with open(os.path.join(self.processed_folder, self.test_file), 'wb') as f:
            torch.save(test_set, f)

        print('Done!')

def
download
(
self
)
:
if
self
.
_check_exists
(
)
:
return
makedir_exist_ok
(
self
.
raw_folder
)
makedir_exist_ok
(
self
.
processed_folder
)
# download files
for
url
in
self
.
urls
:
filename
=
url
.
rpartition
(
'/'
)
[
2
]
file_path
=
os
.
path
.
join
(
self
.
raw_folder
,
filename
)
download_url
(
url
,
root
=
self
.
raw_folder
,
filename
=
filename
,
md5
=
)
self
.
extract_gzip
(
gzip_path
=
file_path
,
remove_finished
=
True
)
# process and save as torch files
print
(
'Processing...'
)
training_set
=
(
read_image_file
(
os
.
path
.
join
(
self
.
raw_folder
,
'train-images-idx3-ubyte'
)
)
,
read_label_file
(
os
.
path
.
join
(
self
.
raw_folder
,
'train-labels-idx1-ubyte'
)
)
)
test_set
=
(
read_image_file
(
os
.
path
.
join
(
self
.
raw_folder
,
't10k-images-idx3-ubyte'
)
)
,
read_label_file
(
os
.
path
.
join
(
self
.
raw_folder
,
't10k-labels-idx1-ubyte'
)
)
)
with
open
(
os
.
path
.
join
(
self
.
processed_folder
,
self
.
training_file
)
,
'wb'
)
as
f
:
torch
.
save
(
training_set
,
f
)
with
open
(
os
.
path
.
join
(
self
.
processed_folder
,
self
.
test_file
)
,
'wb'
)
as
f
:
torch
.
save
(
test_set
,
f
)
print
(
'Done!'
)

Download the MNIST data if it doesn t exist in processed_folder already .


def download(self):
        """Download the EMNIST data if it doesn't exist in processed_folder already."""
        import shutil
        import zipfile

        if self._check_exists():
            return

        makedir_exist_ok(self.raw_folder)
        makedir_exist_ok(self.processed_folder)

        # download files
        filename = self.url.rpartition('/')[2]
        file_path = os.path.join(self.raw_folder, filename)
        download_url(self.url, root=self.raw_folder, filename=filename, md5=None)

        print('Extracting zip archive')
        with zipfile.ZipFile(file_path) as zip_f:
            zip_f.extractall(self.raw_folder)
        os.unlink(file_path)
        gzip_folder = os.path.join(self.raw_folder, 'gzip')
        for gzip_file in os.listdir(gzip_folder):
            if gzip_file.endswith('.gz'):
                self.extract_gzip(gzip_path=os.path.join(gzip_folder, gzip_file))

        # process and save as torch files
        for split in self.splits:
            print('Processing ' + split)
            training_set = (
                read_image_file(os.path.join(gzip_folder, 'emnist-{}-train-images-idx3-ubyte'.format(split))),
                read_label_file(os.path.join(gzip_folder, 'emnist-{}-train-labels-idx1-ubyte'.format(split)))
            )
            test_set = (
                read_image_file(os.path.join(gzip_folder, 'emnist-{}-test-images-idx3-ubyte'.format(split))),
                read_label_file(os.path.join(gzip_folder, 'emnist-{}-test-labels-idx1-ubyte'.format(split)))
            )
            with open(os.path.join(self.processed_folder, self._training_file(split)), 'wb') as f:
                torch.save(training_set, f)
            with open(os.path.join(self.processed_folder, self._test_file(split)), 'wb') as f:
                torch.save(test_set, f)
        shutil.rmtree(gzip_folder)

        print('Done!')

def
download
(
self
)
:
import
shutil
import
zipfile
if
self
.
_check_exists
(
)
:
return
makedir_exist_ok
(
self
.
raw_folder
)
makedir_exist_ok
(
self
.
processed_folder
)
# download files
filename
=
self
.
url
.
rpartition
(
'/'
)
[
2
]
file_path
=
os
.
path
.
join
(
self
.
raw_folder
,
filename
)
download_url
(
self
.
url
,
root
=
self
.
raw_folder
,
filename
=
filename
,
md5
=
)
print
(
'Extracting zip archive'
)
with
zipfile
.
ZipFile
(
file_path
)
as
zip_f
:
zip_f
.
extractall
(
self
.
raw_folder
)
os
.
unlink
(
file_path
)
gzip_folder
=
os
.
path
.
join
(
self
.
raw_folder
,
'gzip'
)
for
gzip_file
in
os
.
listdir
(
gzip_folder
)
:
if
gzip_file
.
endswith
(
'.gz'
)
:
self
.
extract_gzip
(
gzip_path
=
os
.
path
.
join
(
gzip_folder
,
gzip_file
)
)
# process and save as torch files
for
split
in
self
.
splits
:
print
(
'Processing '
+
split
)
training_set
=
(
read_image_file
(
os
.
path
.
join
(
gzip_folder
,
'emnist-{}-train-images-idx3-ubyte'
.
format
(
split
)
)
)
,
read_label_file
(
os
.
path
.
join
(
gzip_folder
,
'emnist-{}-train-labels-idx1-ubyte'
.
format
(
split
)
)
)
)
test_set
=
(
read_image_file
(
os
.
path
.
join
(
gzip_folder
,
'emnist-{}-test-images-idx3-ubyte'
.
format
(
split
)
)
)
,
read_label_file
(
os
.
path
.
join
(
gzip_folder
,
'emnist-{}-test-labels-idx1-ubyte'
.
format
(
split
)
)
)
)
with
open
(
os
.
path
.
join
(
self
.
processed_folder
,
self
.
_training_file
(
split
)
)
,
'wb'
)
as
f
:
torch
.
save
(
training_set
,
f
)
with
open
(
os
.
path
.
join
(
self
.
processed_folder
,
self
.
_test_file
(
split
)
)
,
'wb'
)
as
f
:
torch
.
save
(
test_set
,
f
)
shutil
.
rmtree
(
gzip_folder
)
print
(
'Done!'
)

Download the EMNIST data if it doesn t exist in processed_folder already .


def get_current_theme_name(override=None):
    """Returns theme name.

    Checks in this order:
    1. override
    2. cookies
    3. settings"""

    if override and (override in themes or override == '__common__'):
        return override
    theme_name = request.args.get('theme', request.preferences.get_value('theme'))
    if theme_name not in themes:
        theme_name = default_theme
    return theme_name

def get_current_theme_name ( override = None ) : if override and ( override in themes or override == '__common__' ) : return override theme_name = request . args . get ( 'theme' , request . preferences . get_value ( 'theme' ) ) if theme_name not in themes : theme_name = default_theme return theme_name


def autocompleter():
    """Return autocompleter results"""

    # set blocked engines
    disabled_engines = request.preferences.engines.get_disabled()

    # parse query
    if PY3:
        raw_text_query = RawTextQuery(request.form.get('q', b''), disabled_engines)
    else:
        raw_text_query = RawTextQuery(request.form.get('q', u'').encode('utf-8'), disabled_engines)
    raw_text_query.parse_query()

    # check if search query is set
    if not raw_text_query.getSearchQuery():
        return '', 400

    # run autocompleter
    completer = autocomplete_backends.get(request.preferences.get_value('autocomplete'))

    # parse searx specific autocompleter results like !bang
    raw_results = searx_bang(raw_text_query)

    # normal autocompletion results only appear if max 3 inner results returned
    if len(raw_results) <= 3 and completer:
        # get language from cookie
        language = request.preferences.get_value('language')
        if not language or language == 'all':
            language = 'en'
        else:
            language = language.split('-')[0]
        # run autocompletion
        raw_results.extend(completer(raw_text_query.getSearchQuery(), language))

    # parse results (write :language and !engine back to result string)
    results = []
    for result in raw_results:
        raw_text_query.changeSearchQuery(result)

        # add parsed result
        results.append(raw_text_query.getFullQuery())

    # return autocompleter results
    if request.form.get('format') == 'x-suggestions':
        return Response(json.dumps([raw_text_query.query, results]),
                        mimetype='application/json')

    return Response(json.dumps(results),
                    mimetype='application/json')


def preferences():
    """Render preferences page && save user preferences"""

    # save preferences
    if request.method == 'POST':
        resp = make_response(redirect(urljoin(settings['server']['base_url'], url_for('index'))))
        try:
            request.preferences.parse_form(request.form)
        except ValidationException:
            request.errors.append(gettext('Invalid settings, please edit your preferences'))
            return resp
        return request.preferences.save(resp)

    # render preferences
    image_proxy = request.preferences.get_value('image_proxy')
    lang = request.preferences.get_value('language')
    disabled_engines = request.preferences.engines.get_disabled()
    allowed_plugins = request.preferences.plugins.get_enabled()

    # stats for preferences page
    stats = {}

    for c in categories:
        for e in categories[c]:
            stats[e.name] = {'time': None,
                             'warn_timeout': False,
                             'warn_time': False}
            if e.timeout > settings['outgoing']['request_timeout']:
                stats[e.name]['warn_timeout'] = True
            stats[e.name]['supports_selected_language'] = _is_selected_language_supported(e, request.preferences)

    # get first element [0], the engine time,
    # and then the second element [1] : the time (the first one is the label)
    for engine_stat in get_engines_stats()[0][1]:
        stats[engine_stat.get('name')]['time'] = round(engine_stat.get('avg'), 3)
        if engine_stat.get('avg') > settings['outgoing']['request_timeout']:
            stats[engine_stat.get('name')]['warn_time'] = True
    # end of stats

    return render('preferences.html',
                  locales=settings['locales'],
                  current_locale=get_locale(),
                  image_proxy=image_proxy,
                  engines_by_category=categories,
                  stats=stats,
                  answerers=[{'info': a.self_info(), 'keywords': a.keywords} for a in answerers],
                  disabled_engines=disabled_engines,
                  autocomplete_backends=autocomplete_backends,
                  shortcuts={y: x for x, y in engine_shortcuts.items()},
                  themes=themes,
                  plugins=plugins,
                  doi_resolvers=settings['doi_resolvers'],
                  current_doi_resolver=get_doi_resolver(request.args, request.preferences.get_value('doi_resolver')),
                  allowed_plugins=allowed_plugins,
                  theme=get_current_theme_name(),
                  preferences_url_params=request.preferences.get_as_url_params(),
                  base_url=get_base_url(),
                  preferences=True)


def get_themes(templates_path):
    """Returns available themes list."""
    themes = os.listdir(templates_path)
    if '__common__' in themes:
        themes.remove('__common__')
    return themes

def get_themes ( templates_path ) : themes = os . listdir ( templates_path ) if '__common__' in themes : themes . remove ( '__common__' ) return themes


def searx_bang(full_query):
    '''check if the searchQuery contain a bang, and create fitting autocompleter results'''
    # check if there is a query which can be parsed
    if len(full_query.getSearchQuery()) == 0:
        return []

    results = []

    # check if current query stats with !bang
    first_char = full_query.getSearchQuery()[0]
    if first_char == '!' or first_char == '?':
        if len(full_query.getSearchQuery()) == 1:
            # show some example queries
            # TODO, check if engine is not avaliable
            results.append(first_char + "images")
            results.append(first_char + "wikipedia")
            results.append(first_char + "osm")
        else:
            engine_query = full_query.getSearchQuery()[1:]

            # check if query starts with categorie name
            for categorie in categories:
                if categorie.startswith(engine_query):
                    results.append(first_char + '{categorie}'.format(categorie=categorie))

            # check if query starts with engine name
            for engine in engines:
                if engine.startswith(engine_query.replace('_', ' ')):
                    results.append(first_char + '{engine}'.format(engine=engine.replace(' ', '_')))

            # check if query starts with engine shortcut
            for engine_shortcut in engine_shortcuts:
                if engine_shortcut.startswith(engine_query):
                    results.append(first_char + '{engine_shortcut}'.format(engine_shortcut=engine_shortcut))

    # check if current query stats with :bang
    elif first_char == ':':
        if len(full_query.getSearchQuery()) == 1:
            # show some example queries
            results.append(":en")
            results.append(":en_us")
            results.append(":english")
            results.append(":united_kingdom")
        else:
            engine_query = full_query.getSearchQuery()[1:]

            for lc in language_codes:
                lang_id, lang_name, country, english_name = map(unicode.lower, lc)

                # check if query starts with language-id
                if lang_id.startswith(engine_query):
                    if len(engine_query) <= 2:
                        results.append(u':{lang_id}'.format(lang_id=lang_id.split('-')[0]))
                    else:
                        results.append(u':{lang_id}'.format(lang_id=lang_id))

                # check if query starts with language name
                if lang_name.startswith(engine_query) or english_name.startswith(engine_query):
                    results.append(u':{lang_name}'.format(lang_name=lang_name))

                # check if query starts with country
                if country.startswith(engine_query.replace('_', ' ')):
                    results.append(u':{country}'.format(country=country.replace(' ', '_')))

    # remove duplicates
    result_set = set(results)

    # remove results which are already contained in the query
    for query_part in full_query.query_parts:
        if query_part in result_set:
            result_set.remove(query_part)

    # convert result_set back to list
    return list(result_set)

check if the searchQuery contain a bang and create fitting autocompleter results


def response(resp):
    """remove first and last lines to get only json"""
    json_resp = resp.text[resp.text.find('\n') + 1:resp.text.rfind('\n') - 2]
    results = []
    try:
        conversion_rate = float(json.loads(json_resp)['conversion']['converted-amount'])
    except:
        return results
    answer = '{0} {1} = {2} {3}, 1 {1} ({5}) = {4} {3} ({6})'.format(
        resp.search_params['amount'],
        resp.search_params['from'],
        resp.search_params['amount'] * conversion_rate,
        resp.search_params['to'],
        conversion_rate,
        resp.search_params['from_name'],
        resp.search_params['to_name'],
    )

    url = 'https://duckduckgo.com/js/spice/currency/1/{0}/{1}'.format(
        resp.search_params['from'].upper(), resp.search_params['to'])

    results.append({'answer': answer, 'url': url})

    return results


def custom_gradient(fx, gx, x, fx_gx_manually_stopped=False, name=None):
  """Embeds a custom gradient into a `Tensor`.

  This function works by clever application of `stop_gradient`. I.e., observe
  that:

  ```none
  h(x) = stop_gradient(f(x)) + stop_gradient(g(x)) * (x - stop_gradient(x))
  ```

  is such that `h(x) == stop_gradient(f(x))` and
  `grad[h(x), x] == stop_gradient(g(x)).`

  In addition to scalar-domain/scalar-range functions, this function also
  supports tensor-domain/scalar-range functions.

  Partial Custom Gradient:

  Suppose `h(x) = htilde(x, y)`. Note that `dh/dx = stop(g(x))` but `dh/dy =
  None`. This is because a `Tensor` cannot have only a portion of its gradient
  stopped. To circumvent this issue, one must manually `stop_gradient` the
  relevant portions of `f`, `g`. For example see the unit-test,
  `test_works_correctly_fx_gx_manually_stopped`.

  Args:
    fx: `Tensor`. Output of function evaluated at `x`.
    gx: `Tensor` or list of `Tensor`s. Gradient of function at (each) `x`.
    x: `Tensor` or list of `Tensor`s. Args of evaluation for `f`.
    fx_gx_manually_stopped: Python `bool` indicating that `fx`, `gx` manually
      have `stop_gradient` applied.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    fx: Floating-type `Tensor` equal to `f(x)` but which has gradient
      `stop_gradient(g(x))`.
  """
  def maybe_stop(x):
    if fx_gx_manually_stopped:
      return x
    return tf.stop_gradient(x)

  with tf.compat.v1.name_scope(name, 'custom_gradient', [fx, gx, x]):
    fx = tf.convert_to_tensor(value=fx, name='fx')
    # We don't want to bother eagerly computing `gx` since we may not even need
    # it.
    with tf.control_dependencies([fx]):
      if is_list_like(x):
        x = [identity(x_, name='x') for x_ in x]
      else:
        x = [identity(x, name='x')]

      if is_list_like(gx):
        gx = [identity(gx_, dtype=fx.dtype, name='gx')
              for gx_ in gx]
      else:
        gx = [identity(gx, dtype=fx.dtype, name='gx')]

      override_grad = []
      for x_, gx_ in zip(x, gx):
        # Observe: tf.gradients(f(x), x)[i].shape == x[i].shape
        # thus we check that the user is supplying correct shapes.
        equal_shape = tf.compat.v1.assert_equal(
            tf.shape(input=x_),
            tf.shape(input=gx_),
            message='Each `x` must have the same shape as each `gx`.')
        with tf.control_dependencies([equal_shape]):
          # IEEE754 ensures `(x-x)==0.` and that `0.*x==0.` so we make sure to
          # write the code this way, rather than, e.g.,
          # `sum_x * stop(gx) + stop(fx - sum_x * gx)`.
          # For more discussion regarding the relevant portions of the IEEE754
          # standard, see the StackOverflow question,
          # "Is there a floating point value of x, for which x-x == 0 is false?"
          # http://stackoverflow.com/q/2686644
          zeros_like_x_ = x_ - tf.stop_gradient(x_)
          override_grad.append(
              tf.reduce_sum(input_tensor=maybe_stop(gx_) * zeros_like_x_))
      override_grad = sum(override_grad)
      override_grad /= tf.cast(tf.size(input=fx), dtype=fx.dtype.base_dtype)

      # Proof of correctness:
      #
      #  f(x) = x * stop[gx] + stop[fx - x * gx]
      #       = stop[fx]
      #
      #  g(x) = grad[fx]
      #       = stop[gx] + grad[stop[fx - x * gx]]
      #       = stop[gx] + 0
      #
      # Notice that when x is zero it still works:
      # grad[x * stop(gx) + stop(fx - x * gx)] = 1 * stop[gx] + 0 = stop[gx]
      #
      # The proof is similar for the tensor-domain case, except that we
      # `reduce_sum` the `stop[gx] * (x - stop[x])` then rescale by
      # `tf.size(fx)` since this reduced version is broadcast to `fx`.
      return maybe_stop(fx) + override_grad

#       = stop[gx] + grad[stop[fx - x * gx]]

Embeds a custom gradient into a Tensor .


def mvn(*args, **kwargs):
  """Convenience function to efficiently construct a MultivariateNormalDiag."""
  # Faster than using `tfd.MultivariateNormalDiag`.
  return tfd.Independent(tfd.Normal(*args, **kwargs),
                         reinterpreted_batch_ndims=1)

Convenience function to efficiently construct a MultivariateNormalDiag .


def eight_schools_joint_log_prob(
    treatment_effects, treatment_stddevs,
    avg_effect, avg_stddev, school_effects_standard):
  """Eight-schools joint log-prob."""
  rv_avg_effect = tfd.Normal(loc=0., scale=10.)
  rv_avg_stddev = tfd.Normal(loc=5., scale=1.)
  rv_school_effects_standard = mvn(
      loc=tf.zeros_like(school_effects_standard),
      scale=tf.ones_like(school_effects_standard))
  rv_treatment_effects = mvn(
      loc=(avg_effect + tf.exp(avg_stddev) * school_effects_standard),
      scale=treatment_stddevs)
  return (
      rv_avg_effect.log_prob(avg_effect) +
      rv_avg_stddev.log_prob(avg_stddev) +
      rv_school_effects_standard.log_prob(school_effects_standard) +
      rv_treatment_effects.log_prob(treatment_effects))

def eight_schools_joint_log_prob ( treatment_effects , treatment_stddevs , avg_effect , avg_stddev , school_effects_standard ) : rv_avg_effect = tfd . Normal ( loc = 0. , scale = 10. ) rv_avg_stddev = tfd . Normal ( loc = 5. , scale = 1. ) rv_school_effects_standard = mvn ( loc = tf . zeros_like ( school_effects_standard ) , scale = tf . ones_like ( school_effects_standard ) ) rv_treatment_effects = mvn ( loc = ( avg_effect + tf . exp ( avg_stddev ) * school_effects_standard ) , scale = treatment_stddevs ) return ( rv_avg_effect . log_prob ( avg_effect ) + rv_avg_stddev . log_prob ( avg_stddev ) + rv_school_effects_standard . log_prob ( school_effects_standard ) + rv_treatment_effects . log_prob ( treatment_effects ) )


def benchmark_eight_schools_hmc(
    num_results=int(5e3),
    num_burnin_steps=int(3e3),
    num_leapfrog_steps=3,
    step_size=0.4):
  """Runs HMC on the eight-schools unnormalized posterior."""

  num_schools = 8
  treatment_effects = tf.constant(
      [28, 8, -3, 7, -1, 1, 18, 12],
      dtype=np.float32,
      name='treatment_effects')
  treatment_stddevs = tf.constant(
      [15, 10, 16, 11, 9, 11, 10, 18],
      dtype=np.float32,
      name='treatment_stddevs')

  def unnormalized_posterior_log_prob(
      avg_effect, avg_stddev, school_effects_standard):
    """Eight-schools unnormalized log posterior."""
    return eight_schools_joint_log_prob(
        treatment_effects, treatment_stddevs,
        avg_effect, avg_stddev, school_effects_standard)

  if tf.executing_eagerly():
    sample_chain = tf.function(tfp.mcmc.sample_chain)
  else:
    sample_chain = tfp.mcmc.sample_chain

  def computation():
    """The benchmark computation."""
    _, kernel_results = sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=(
            tf.zeros([], name='init_avg_effect'),
            tf.zeros([], name='init_avg_stddev'),
            tf.ones([num_schools], name='init_school_effects_standard'),
        ),
        kernel=tfp.mcmc.HamiltonianMonteCarlo(
            target_log_prob_fn=unnormalized_posterior_log_prob,
            step_size=step_size,
            num_leapfrog_steps=num_leapfrog_steps))

    return kernel_results.is_accepted

  # Let's force evaluation of graph to ensure build time is not part of our time
  # trial.
  is_accepted_tensor = computation()
  if not tf.executing_eagerly():
    session = tf.compat.v1.Session()
    session.run(is_accepted_tensor)

  start_time = time.time()
  if tf.executing_eagerly():
    is_accepted = computation()
  else:
    is_accepted = session.run(is_accepted_tensor)
  wall_time = time.time() - start_time

  num_accepted = np.sum(is_accepted)
  acceptance_rate = np.float32(num_accepted) / np.float32(num_results)

  return dict(
      iters=(num_results + num_burnin_steps) * num_leapfrog_steps,
      extras={'acceptance_rate': acceptance_rate},
      wall_time=wall_time)

Runs HMC on the eight - schools unnormalized posterior .


def expand_docstring(**kwargs):
  """Decorator to programmatically expand the docstring.

  Args:
    **kwargs: Keyword arguments to set. For each key-value pair `k` and `v`,
      the key is found as `${k}` in the docstring and replaced with `v`.

  Returns:
    Decorated function.
  """
  def _fn_wrapped(fn):
    """Original function with modified `__doc__` attribute."""
    doc = inspect.cleandoc(fn.__doc__)
    for k, v in six.iteritems(kwargs):
      # Capture each ${k} reference to replace with v.
      # We wrap the replacement in a function so no backslash escapes
      # are processed.
      pattern = r'\$\{' + str(k) + r'\}'
      doc = re.sub(pattern, lambda match: v, doc)  # pylint: disable=cell-var-from-loop
    fn.__doc__ = doc
    return fn
  return _fn_wrapped

Decorator to programmatically expand the docstring .


def _simple_name(distribution):
  """Infer the original name passed into a distribution constructor.

  Distributions typically follow the pattern of
  with.name_scope(name) as name:
    super(name=name)
  so we attempt to reverse the name-scope transformation to allow
  addressing of RVs by the distribution's original, user-visible
  name kwarg.

  Args:
    distribution: a tfd.Distribution instance.
  Returns:
    simple_name: the original name passed into the Distribution.

  #### Example

  ```
  d1 = tfd.Normal(0., 1., name='x') # d1.name = 'x/'
  d2 = tfd.Normal(0., 1., name='x') # d2.name = 'x_2/'
  _simple_name(d2) # returns 'x'

  ```

  """
  simple_name = distribution.name

  # turn 'scope/x/' into 'x'
  if simple_name.endswith('/'):
    simple_name = simple_name.split('/')[-2]

  # turn 'x_3' into 'x'
  parts = simple_name.split('_')
  if parts[-1].isdigit():
    simple_name = '_'.join(parts[:-1])

  return simple_name

def
_simple_name
(
distribution
)
:
simple_name
=
distribution
.
name
# turn 'scope/x/' into 'x'
if
simple_name
.
endswith
(
'/'
)
:
simple_name
=
simple_name
.
split
(
'/'
)
[
2
]
# turn 'x_3' into 'x'
parts
=
simple_name
.
split
(
'_'
)
if
parts
[
1
]
.
isdigit
(
)
:
simple_name
=
'_'
.
join
(
parts
[
:
1
]
)
return
simple_name

Infer the original name passed into a distribution constructor .


def _build_custom_rv(distribution, sample_shape, value, name):
  """RandomVariable constructor with a dummy name argument."""
  # Program transformations (e.g., `make_log_joint_fn`) assume that
  # the traced constructor has `name` and `value` kwargs, enabling
  # them to override the value of an RV according to its name.
  # User-defined RVs inherit their name from the provided
  # distribution; this helper method exposes the name as a dummy kwarg
  # so that it's visible to program transformations.
  del name  # unused
  return RandomVariable(distribution=distribution,
                        sample_shape=sample_shape,
                        value=value)

RandomVariable constructor with a dummy name argument .


def as_random_variable(distribution,
                       sample_shape=(),
                       value=None):
  """Wrap an existing distribution as a traceable random variable.

  This enables the use of custom or user-provided distributions in
  Edward models. Unlike a bare `RandomVariable` object, this method
  wraps the constructor so it is included in the Edward trace and its
  values can be properly intercepted and overridden.

  Where possible, you should prefer the built-in constructors
  (`ed.Normal`, etc); these simultaneously construct a Distribution
  and a RandomVariable object so that the distribution parameters
  themselves may be intercepted and overridden. RVs constructed via
  `as_random_variable()` have a fixed distribution and may not support
  program transformations (e.g, conjugate marginalization) that rely
  on overriding distribution parameters.

  Args:
    distribution: tfd.Distribution governing the distribution of the random
      variable, such as sampling and log-probabilities.
    sample_shape: tf.TensorShape of samples to draw from the random variable.
      Default is `()` corresponding to a single sample.
    value: Fixed tf.Tensor to associate with random variable. Must have shape
      `sample_shape + distribution.batch_shape + distribution.event_shape`.
      Default is to sample from random variable according to `sample_shape`.

  Returns:
    rv: a `RandomVariable` wrapping the provided distribution.

  #### Example

  ```python
  from tensorflow_probability import distributions as tfd
  from tensorflow_probability import edward2 as ed

  def model():
    # equivalent to ed.Normal(0., 1., name='x')
    return ed.as_random_variable(tfd.Normal(0., 1., name='x'))

  log_joint = ed.make_log_joint_fn(model)
  output = log_joint(x=2.)
  ```
  """

  return _build_custom_rv(distribution=distribution,
                          sample_shape=sample_shape,
                          value=value,
                          name=_simple_name(distribution))

def as_random_variable ( distribution , sample_shape = ( ) , value = None ) : return _build_custom_rv ( distribution = distribution , sample_shape = sample_shape , value = value , name = _simple_name ( distribution ) )

Wrap an existing distribution as a traceable random variable .


def _make_random_variable(distribution_cls):
  """Factory function to make random variable given distribution class."""

  @interceptable
  @functools.wraps(distribution_cls, assigned=('__module__', '__name__'))
  @docstring_util.expand_docstring(
      cls=distribution_cls.__name__,
      doc=inspect.cleandoc(distribution_cls.__init__.__doc__ or ''))
  def func(*args, **kwargs):
    # pylint: disable=g-doc-args
    """Create a random variable for ${cls}.

    See ${cls} for more details.

    Returns:
      RandomVariable.

    #### Original Docstring for Distribution

    ${doc}
    """
    # pylint: enable=g-doc-args
    sample_shape = kwargs.pop('sample_shape', ())
    value = kwargs.pop('value', None)
    return RandomVariable(distribution=distribution_cls(*args, **kwargs),
                          sample_shape=sample_shape,
                          value=value)
  return func

Factory function to make random variable given distribution class .


def one_step_predictive(model, observed_time_series, parameter_samples):
  """Compute one-step-ahead predictive distributions for all timesteps.

  Given samples from the posterior over parameters, return the predictive
  distribution over observations at each time `T`, given observations up
  through time `T-1`.

  Args:
    model: An instance of `StructuralTimeSeries` representing a
      time-series model. This represents a joint distribution over
      time-series and their parameters with batch shape `[b1, ..., bN]`.
    observed_time_series: `float` `Tensor` of shape
      `concat([sample_shape, model.batch_shape, [num_timesteps, 1]]) where
      `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
      dimension may (optionally) be omitted if `num_timesteps > 1`. May
      optionally be an instance of `tfp.sts.MaskedTimeSeries` including a
      mask `Tensor` to encode the locations of missing observations.
    parameter_samples: Python `list` of `Tensors` representing posterior samples
      of model parameters, with shapes `[concat([[num_posterior_draws],
      param.prior.batch_shape, param.prior.event_shape]) for param in
      model.parameters]`. This may optionally also be a map (Python `dict`) of
      parameter names to `Tensor` values.

  Returns:
    forecast_dist: a `tfd.MixtureSameFamily` instance with event shape
      [num_timesteps] and
      batch shape `concat([sample_shape, model.batch_shape])`, with
      `num_posterior_draws` mixture components. The `t`th step represents the
      forecast distribution `p(observed_time_series[t] |
      observed_time_series[0:t-1], parameter_samples)`.

  #### Examples

  Suppose we've built a model and fit it to data using HMC:

  ```python
    day_of_week = tfp.sts.Seasonal(
        num_seasons=7,
        observed_time_series=observed_time_series,
        name='day_of_week')
    local_linear_trend = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series,
        name='local_linear_trend')
    model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                        observed_time_series=observed_time_series)

    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
  ```

  Passing the posterior samples into `one_step_predictive`, we construct a
  one-step-ahead predictive distribution:

  ```python
    one_step_predictive_dist = tfp.sts.one_step_predictive(
      model, observed_time_series, parameter_samples=samples)

    predictive_means = one_step_predictive_dist.mean()
    predictive_scales = one_step_predictive_dist.stddev()
  ```

  If using variational inference instead of HMC, we'd construct a forecast using
  samples from the variational posterior:

  ```python
    (variational_loss,
     variational_distributions) = tfp.sts.build_factored_variational_loss(
       model=model, observed_time_series=observed_time_series)

    # OMITTED: take steps to optimize variational loss

    samples = {k: q.sample(30) for (k, q) in variational_distributions.items()}
    one_step_predictive_dist = tfp.sts.one_step_predictive(
      model, observed_time_series, parameter_samples=samples)
  ```

  We can visualize the forecast by plotting:

  ```python
    from matplotlib import pylab as plt
    def plot_one_step_predictive(observed_time_series,
                                 forecast_mean,
                                 forecast_scale):
      plt.figure(figsize=(12, 6))
      num_timesteps = forecast_mean.shape[-1]
      c1, c2 = (0.12, 0.47, 0.71), (1.0, 0.5, 0.05)
      plt.plot(observed_time_series, label="observed time series", color=c1)
      plt.plot(forecast_mean, label="one-step prediction", color=c2)
      plt.fill_between(np.arange(num_timesteps),
                       forecast_mean - 2 * forecast_scale,
                       forecast_mean + 2 * forecast_scale,
                       alpha=0.1, color=c2)
      plt.legend()

    plot_one_step_predictive(observed_time_series,
                             forecast_mean=predictive_means,
                             forecast_scale=predictive_scales)
  ```

  To detect anomalous timesteps, we check whether the observed value at each
  step is within a 95% predictive interval, i.e., two standard deviations from
  the mean:

  ```python
    z_scores = ((observed_time_series[..., 1:] - predictive_means[..., :-1])
                 / predictive_scales[..., :-1])
    anomalous_timesteps = tf.boolean_mask(
        tf.range(1, num_timesteps),
        tf.abs(z_scores) > 2.0)
  ```

  """

  with tf.compat.v1.name_scope(
      'one_step_predictive', values=[observed_time_series, parameter_samples]):

    [
        observed_time_series,
        is_missing
    ] = sts_util.canonicalize_observed_time_series_with_mask(
        observed_time_series)

    # Run filtering over the training timesteps to extract the
    # predictive means and variances.
    num_timesteps = dist_util.prefer_static_value(
        tf.shape(input=observed_time_series))[-2]
    lgssm = model.make_state_space_model(
        num_timesteps=num_timesteps, param_vals=parameter_samples)
    (_, _, _, _, _, observation_means, observation_covs
    ) = lgssm.forward_filter(observed_time_series, mask=is_missing)

    # Squeeze dims to convert from LGSSM's event shape `[num_timesteps, 1]`
    # to a scalar time series.
    return sts_util.mix_over_posterior_draws(
        means=observation_means[..., 0],
        variances=observation_covs[..., 0, 0])

Compute one - step - ahead predictive distributions for all timesteps .


def forecast(model,
             observed_time_series,
             parameter_samples,
             num_steps_forecast):
  """Construct predictive distribution over future observations.

  Given samples from the posterior over parameters, return the predictive
  distribution over future observations for num_steps_forecast timesteps.

  Args:
    model: An instance of `StructuralTimeSeries` representing a
      time-series model. This represents a joint distribution over
      time-series and their parameters with batch shape `[b1, ..., bN]`.
    observed_time_series: `float` `Tensor` of shape
      `concat([sample_shape, model.batch_shape, [num_timesteps, 1]])` where
      `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
      dimension may (optionally) be omitted if `num_timesteps > 1`. May
      optionally be an instance of `tfp.sts.MaskedTimeSeries` including a
      mask `Tensor` to encode the locations of missing observations.
    parameter_samples: Python `list` of `Tensors` representing posterior samples
      of model parameters, with shapes `[concat([[num_posterior_draws],
      param.prior.batch_shape, param.prior.event_shape]) for param in
      model.parameters]`. This may optionally also be a map (Python `dict`) of
      parameter names to `Tensor` values.
    num_steps_forecast: scalar `int` `Tensor` number of steps to forecast.

  Returns:
    forecast_dist: a `tfd.MixtureSameFamily` instance with event shape
      [num_steps_forecast, 1] and batch shape
      `concat([sample_shape, model.batch_shape])`, with `num_posterior_draws`
      mixture components.

  #### Examples

  Suppose we've built a model and fit it to data using HMC:

  ```python
    day_of_week = tfp.sts.Seasonal(
        num_seasons=7,
        observed_time_series=observed_time_series,
        name='day_of_week')
    local_linear_trend = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series,
        name='local_linear_trend')
    model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                        observed_time_series=observed_time_series)

    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
  ```

  Passing the posterior samples into `forecast`, we construct a forecast
  distribution:

  ```python
    forecast_dist = tfp.sts.forecast(model, observed_time_series,
                                     parameter_samples=samples,
                                     num_steps_forecast=50)

    forecast_mean = forecast_dist.mean()[..., 0]  # shape: [50]
    forecast_scale = forecast_dist.stddev()[..., 0]  # shape: [50]
    forecast_samples = forecast_dist.sample(10)[..., 0]  # shape: [10, 50]
  ```

  If using variational inference instead of HMC, we'd construct a forecast using
  samples from the variational posterior:

  ```python
    (variational_loss,
     variational_distributions) = tfp.sts.build_factored_variational_loss(
       model=model, observed_time_series=observed_time_series)

    # OMITTED: take steps to optimize variational loss

    samples = {k: q.sample(30) for (k, q) in variational_distributions.items()}
    forecast_dist = tfp.sts.forecast(model, observed_time_series,
                                         parameter_samples=samples,
                                         num_steps_forecast=50)
  ```

  We can visualize the forecast by plotting:

  ```python
    from matplotlib import pylab as plt
    def plot_forecast(observed_time_series,
                      forecast_mean,
                      forecast_scale,
                      forecast_samples):
      plt.figure(figsize=(12, 6))

      num_steps = observed_time_series.shape[-1]
      num_steps_forecast = forecast_mean.shape[-1]
      num_steps_train = num_steps - num_steps_forecast

      c1, c2 = (0.12, 0.47, 0.71), (1.0, 0.5, 0.05)
      plt.plot(np.arange(num_steps), observed_time_series,
               lw=2, color=c1, label='ground truth')

      forecast_steps = np.arange(num_steps_train,
                       num_steps_train+num_steps_forecast)
      plt.plot(forecast_steps, forecast_samples.T, lw=1, color=c2, alpha=0.1)
      plt.plot(forecast_steps, forecast_mean, lw=2, ls='--', color=c2,
               label='forecast')
      plt.fill_between(forecast_steps,
                       forecast_mean - 2 * forecast_scale,
                       forecast_mean + 2 * forecast_scale, color=c2, alpha=0.2)

      plt.xlim([0, num_steps])
      plt.legend()

    plot_forecast(observed_time_series,
                  forecast_mean=forecast_mean,
                  forecast_scale=forecast_scale,
                  forecast_samples=forecast_samples)
  ```

  """

  with tf.compat.v1.name_scope(
      'forecast',
      values=[observed_time_series, parameter_samples, num_steps_forecast]):
    [
        observed_time_series,
        mask
    ] = sts_util.canonicalize_observed_time_series_with_mask(
        observed_time_series)

    # Run filtering over the observed timesteps to extract the
    # latent state posterior at timestep T+1 (i.e., the final
    # filtering distribution, pushed through the transition model).
    # This is the prior for the forecast model ("today's prior
    # is yesterday's posterior").
    num_observed_steps = dist_util.prefer_static_value(
        tf.shape(input=observed_time_series))[-2]
    observed_data_ssm = model.make_state_space_model(
        num_timesteps=num_observed_steps, param_vals=parameter_samples)
    (_, _, _, predictive_means, predictive_covs, _, _
    ) = observed_data_ssm.forward_filter(observed_time_series, mask=mask)

    # Build a batch of state-space models over the forecast period. Because
    # we'll use MixtureSameFamily to mix over the posterior draws, we need to
    # do some shenanigans to move the `[num_posterior_draws]` batch dimension
    # from the leftmost to the rightmost side of the model's batch shape.
    # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
    # arbitrary axis, and eliminate `move_dimension` calls here.
    parameter_samples = model._canonicalize_param_vals_as_map(parameter_samples)  # pylint: disable=protected-access
    parameter_samples_with_reordered_batch_dimension = {
        param.name: dist_util.move_dimension(
            parameter_samples[param.name],
            0, -(1 + _prefer_static_event_ndims(param.prior)))
        for param in model.parameters}
    forecast_prior = tfd.MultivariateNormalFullCovariance(
        loc=dist_util.move_dimension(predictive_means[..., -1, :], 0, -2),
        covariance_matrix=dist_util.move_dimension(
            predictive_covs[..., -1, :, :], 0, -3))

    # Ugly hack: because we moved `num_posterior_draws` to the trailing (rather
    # than leading) dimension of parameters, the parameter batch shapes no
    # longer broadcast against the `constant_offset` attribute used in `sts.Sum`
    # models. We fix this by manually adding an extra broadcasting dim to
    # `constant_offset` if present.
    # The root cause of this hack is that we mucked with param dimensions above
    # and are now passing params that are 'invalid' in the sense that they don't
    # match the shapes of the model's param priors. The fix (as above) will be
    # to update MixtureSameFamily so we can avoid changing param dimensions
    # altogether.
    # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
    # arbitrary axis, and eliminate this hack.
    kwargs = {}
    if hasattr(model, 'constant_offset'):
      kwargs['constant_offset'] = tf.convert_to_tensor(
          value=model.constant_offset,
          dtype=forecast_prior.dtype)[..., tf.newaxis]

    # We assume that any STS model that has a `constant_offset` attribute
    # will allow it to be overridden as a kwarg. This is currently just
    # `sts.Sum`.
    # TODO(b/120245392): when kwargs hack is removed, switch back to calling
    # the public version of `_make_state_space_model`.
    forecast_ssm = model._make_state_space_model(  # pylint: disable=protected-access
        num_timesteps=num_steps_forecast,
        param_map=parameter_samples_with_reordered_batch_dimension,
        initial_state_prior=forecast_prior,
        initial_step=num_observed_steps,
        **kwargs)

    num_posterior_draws = dist_util.prefer_static_value(
        forecast_ssm.batch_shape_tensor())[-1]
    return tfd.MixtureSameFamily(
        mixture_distribution=tfd.Categorical(
            logits=tf.zeros([num_posterior_draws], dtype=forecast_ssm.dtype)),
        components_distribution=forecast_ssm)

Construct predictive distribution over future observations .


def _max_mask_non_finite(x, axis=-1, keepdims=False, mask=0):
  """Returns `max` or `mask` if `max` is not finite."""
  m = np.max(x, axis=_astuple(axis), keepdims=keepdims)
  needs_masking = ~np.isfinite(m)
  if needs_masking.ndim > 0:
    m[needs_masking] = mask
  elif needs_masking:
    m = mask
  return m

def _max_mask_non_finite ( x , axis = - 1 , keepdims = False , mask = 0 ) : m = np . max ( x , axis = _astuple ( axis ) , keepdims = keepdims ) needs_masking = ~ np . isfinite ( m ) if needs_masking . ndim > 0 : m [ needs_masking ] = mask elif needs_masking : m = mask return m

Returns max or mask if max is not finite .


def assert_finite(x, data=None, summarize=None, message=None, name=None):
  """Assert all elements of `x` are finite.

  Args:
    x:  Numeric `Tensor`.
    data:  The tensors to print out if the condition is False.  Defaults to
      error message and first few entries of `x`.
    summarize: Print this many entries of each tensor.
    message: A string to prefix to the default message.
    name: A name for this operation (optional).
      Defaults to "assert_finite".

  Returns:
    Op raising `InvalidArgumentError` unless `x` has specified rank or lower.
    If static checks determine `x` has correct rank, a `no_op` is returned.

  Raises:
    ValueError:  If static checks determine `x` has wrong rank.
  """
  with tf.compat.v2.name_scope(name or 'assert_finite'):
    x_ = tf.get_static_value(x)
    if x_ is not None:
      if ~np.all(np.isfinite(x_)):
        raise ValueError(message)
      return x
    assertion = tf.compat.v1.assert_equal(
        tf.math.is_finite(x), tf.ones_like(x, tf.bool),
        data=data, summarize=summarize, message=message)
    with tf.control_dependencies([assertion]):
      return tf.identity(x)

def assert_finite ( x , data = None , summarize = None , message = None , name = None ) : with tf . compat . v2 . name_scope ( name or 'assert_finite' ) : x_ = tf . get_static_value ( x ) if x_ is not None : if ~ np . all ( np . isfinite ( x_ ) ) : raise ValueError ( message ) return x assertion = tf . compat . v1 . assert_equal ( tf . math . is_finite ( x ) , tf . ones_like ( x , tf . bool ) , data = data , summarize = summarize , message = message ) with tf . control_dependencies ( [ assertion ] ) : return tf . identity ( x )


def assert_rank_at_most(x, rank, data=None, summarize=None, message=None,
                        name=None):
  """Assert `x` has rank equal to `rank` or smaller.

  Example of adding a dependency to an operation:

  ```python
  with tf.control_dependencies([tf.assert_rank_at_most(x, 2)]):
    output = tf.reduce_sum(x)
  ```

  Args:
    x:  Numeric `Tensor`.
    rank:  Scalar `Tensor`.
    data:  The tensors to print out if the condition is False.  Defaults to
      error message and first few entries of `x`.
    summarize: Print this many entries of each tensor.
    message: A string to prefix to the default message.
    name: A name for this operation (optional).
      Defaults to "assert_rank_at_most".

  Returns:
    Op raising `InvalidArgumentError` unless `x` has specified rank or lower.
    If static checks determine `x` has correct rank, a `no_op` is returned.

  Raises:
    ValueError:  If static checks determine `x` has wrong rank.
  """
  with tf.compat.v2.name_scope(name or 'assert_rank_at_most'):
    return tf.compat.v1.assert_less_equal(
        tf.rank(x), rank, data=data, summarize=summarize, message=message)

def assert_rank_at_most ( x , rank , data = None , summarize = None , message = None , name = None ) : with tf . compat . v2 . name_scope ( name or 'assert_rank_at_most' ) : return tf . compat . v1 . assert_less_equal ( tf . rank ( x ) , rank , data = data , summarize = summarize , message = message )

Assert x has rank equal to rank or smaller .


def _event_size(event_shape, name=None):
  """Computes the number of elements in a tensor with shape `event_shape`.

  Args:
    event_shape: A tensor shape.
    name: The name to use for the tensor op to compute the number of elements
      (if such an op needs to be created).

  Returns:
    event_size: The number of elements in `tensor_shape`.  Returns a numpy int
    when the number of elements can be computed immediately.  Otherwise, returns
    a scalar tensor.
  """
  with tf.compat.v1.name_scope(name, 'event_size', [event_shape]):
    event_shape = tf.convert_to_tensor(
        value=event_shape, dtype=tf.int32, name='event_shape')

    event_shape_const = tf.get_static_value(event_shape)
    if event_shape_const is not None:
      return np.prod(event_shape_const)
    else:
      return tf.reduce_prod(input_tensor=event_shape)

def _event_size ( event_shape , name = None ) : with tf . compat . v1 . name_scope ( name , 'event_size' , [ event_shape ] ) : event_shape = tf . convert_to_tensor ( value = event_shape , dtype = tf . int32 , name = 'event_shape' ) event_shape_const = tf . get_static_value ( event_shape ) if event_shape_const is not None : return np . prod ( event_shape_const ) else : return tf . reduce_prod ( input_tensor = event_shape )

Computes the number of elements in a tensor with shape event_shape .


def _eval_all_one_hot(fn, dist, name=None):
  """OneHotCategorical helper computing probs, cdf, etc over its support."""
  with tf.compat.v1.name_scope(name, 'eval_all_one_hot'):
    event_size = dist.event_shape_tensor()[-1]
    batch_ndims = tf.size(input=dist.batch_shape_tensor())
    # Reshape `eye(d)` to: `[d] + [1]*batch_ndims + [d]`.
    x = tf.reshape(
        tf.eye(event_size, dtype=dist.dtype),
        shape=tf.pad(
            tensor=tf.ones(batch_ndims, tf.int32),
            paddings=[[1, 1]],
            constant_values=event_size))
    # Compute `fn(x)` then cyclically left-transpose one dim.
    perm = tf.pad(tensor=tf.range(1, batch_ndims + 1), paddings=[[0, 1]])
    return tf.transpose(a=fn(dist, x), perm=perm)

OneHotCategorical helper computing probs cdf etc over its support .


def _get_convert_to_tensor_fn(identifier):
  """Return a convert-to-tensor func, given a name, config, callable, etc."""
  if identifier is None:
    return None

  if isinstance(identifier, six.string_types):
    identifier = str(identifier)
    return _deserialize(identifier)

  if isinstance(identifier, dict):
    return _deserialize(identifier)

  if isinstance(identifier, property):
    identifier = identifier.fget
  if callable(identifier):
    return identifier

  raise ValueError('Could not interpret '
                   'convert-to-tensor function identifier:', identifier)

Return a convert - to - tensor func given a name config callable etc .


def params_size(num_components, component_params_size, name=None):
    """Number of `params` needed to create a `MixtureSameFamily` distribution.

    Arguments:
      num_components: Number of component distributions in the mixture
        distribution.
      component_params_size: Number of parameters needed to create a single
        component distribution.
      name: The name to use for the op to compute the number of parameters
        (if such an op needs to be created).

    Returns:
     params_size: The number of parameters needed to create the mixture
       distribution.
    """
    with tf.compat.v1.name_scope(name, 'MixtureSameFamily_params_size',
                                 [num_components, component_params_size]):
      num_components = tf.convert_to_tensor(
          value=num_components, name='num_components', dtype_hint=tf.int32)
      component_params_size = tf.convert_to_tensor(
          value=component_params_size, name='component_params_size')

      num_components = dist_util.prefer_static_value(num_components)
      component_params_size = dist_util.prefer_static_value(
          component_params_size)

      return num_components + num_components * component_params_size

def params_size ( num_components , component_params_size , name = None ) : with tf . compat . v1 . name_scope ( name , 'MixtureSameFamily_params_size' , [ num_components , component_params_size ] ) : num_components = tf . convert_to_tensor ( value = num_components , name = 'num_components' , dtype_hint = tf . int32 ) component_params_size = tf . convert_to_tensor ( value = component_params_size , name = 'component_params_size' ) num_components = dist_util . prefer_static_value ( num_components ) component_params_size = dist_util . prefer_static_value ( component_params_size ) return num_components + num_components * component_params_size

Number of params needed to create a MixtureSameFamily distribution .


def get_next_interceptor():
  """Yields the top-most interceptor on the thread-local interceptor stack.

  Operations may be intercepted by multiple nested interceptors. Once reached,
  an operation can be forwarded through nested interceptors until resolved.
  To allow for nesting, implement interceptors by re-wrapping their first
  argument (`f`) as an `interceptable`. To avoid nesting, manipulate the
  computation without using `interceptable`.

  This function allows for nesting by manipulating the thread-local interceptor
  stack, so that operations are intercepted in the order of interceptor nesting.

  #### Examples

  ```python
  from tensorflow_probability import edward2 as ed

  def model():
    x = ed.Normal(loc=0., scale=1., name="x")
    y = ed.Normal(loc=x, scale=1., name="y")
    return x + y

  def double(f, *args, **kwargs):
    return 2. * interceptable(f)(*args, **kwargs)

  def set_y(f, *args, **kwargs):
    if kwargs.get("name") == "y":
      kwargs["value"] = 0.42
    return interceptable(f)(*args, **kwargs)

  with interception(double):
    with interception(set_y):
      z = model()
  ```

  This will firstly put `double` on the stack, and then `set_y`,
  resulting in the stack:
  (TOP) set_y -> double -> apply (BOTTOM)

  The execution of `model` is then (top lines are current stack state):
  1) (TOP) set_y -> double -> apply (BOTTOM);
  `ed.Normal(0., 1., "x")` is intercepted by `set_y`, and as the name is not "y"
  the operation is simply forwarded to the next interceptor on the stack.

  2) (TOP) double -> apply (BOTTOM);
  `ed.Normal(0., 1., "x")` is intercepted by `double`, to produce
  `2*ed.Normal(0., 1., "x")`, with the operation being forwarded down the stack.

  3) (TOP) apply (BOTTOM);
  `ed.Normal(0., 1., "x")` is intercepted by `apply`, which simply calls the
  constructor.

  (At this point, the nested calls to `get_next_interceptor()`, produced by
  forwarding operations, exit, and the current stack is again:
  (TOP) set_y -> double -> apply (BOTTOM))

  4) (TOP) set_y -> double -> apply (BOTTOM);
  `ed.Normal(0., 1., "y")` is intercepted by `set_y`,
  the value of `y` is set to 0.42 and the operation is forwarded down the stack.

  5) (TOP) double -> apply (BOTTOM);
  `ed.Normal(0., 1., "y")` is intercepted by `double`, to produce
  `2*ed.Normal(0., 1., "y")`, with the operation being forwarded down the stack.

  6) (TOP) apply (BOTTOM);
  `ed.Normal(0., 1., "y")` is intercepted by `apply`, which simply calls the
  constructor.

  The final values for `x` and `y` inside of `model()` are tensors where `x` is
  a random draw from Normal(0., 1.) doubled, and `y` is a constant 0.84, thus
  z = 2 * Normal(0., 1.) + 0.84.
  """
  try:
    interceptor = _interceptor_stack.stack.pop()
    yield interceptor
  finally:
    _interceptor_stack.stack.append(interceptor)

def get_next_interceptor ( ) : try : interceptor = _interceptor_stack . stack . pop ( ) yield interceptor finally : _interceptor_stack . stack . append ( interceptor )

Yields the top - most interceptor on the thread - local interceptor stack .


def interceptable(func):
  """Decorator that wraps `func` so that its execution is intercepted.

  The wrapper passes `func` to the interceptor for the current thread.

  If there is no next interceptor, we perform an "immediate" call to `func`.
  That is, `func` terminates without forwarding its execution to another
  interceptor.

  Args:
    func: Function to wrap.

  Returns:
    The decorated function.
  """
  @functools.wraps(func)
  def func_wrapped(*args, **kwargs):
    with get_next_interceptor() as interceptor:
      return interceptor(func, *args, **kwargs)

  return func_wrapped

def interceptable ( func ) : @ functools . wraps ( func ) def func_wrapped ( * args , * * kwargs ) : with get_next_interceptor ( ) as interceptor : return interceptor ( func , * args , * * kwargs ) return func_wrapped

Decorator that wraps func so that its execution is intercepted .


def tape():
  """Context manager for recording interceptable executions onto a tape.

  Similar to `tf.GradientTape`, operations are recorded if they are executed
  within this context manager. In addition, the operation must be registered
  (wrapped) as `ed.interceptable`.

  Yields:
    tape: OrderedDict where operations are recorded in sequence. Keys are
      the `name` keyword argument to the operation (typically, a random
      variable's `name`) and values are the corresponding output of the
      operation. If the operation has no name, it is not recorded.

  #### Examples

  ```python
  from tensorflow_probability import edward2 as ed

  def probabilistic_matrix_factorization():
    users = ed.Normal(0., 1., sample_shape=[5000, 128], name="users")
    items = ed.Normal(0., 1., sample_shape=[7500, 128], name="items")
    ratings = ed.Normal(loc=tf.matmul(users, items, transpose_b=True),
                        scale=0.1,
                        name="ratings")
    return ratings

  with ed.tape() as model_tape:
    ratings = probabilistic_matrix_factorization()

  assert model_tape["users"].shape == (5000, 128)
  assert model_tape["items"].shape == (7500, 128)
  assert model_tape["ratings"] == ratings
  ```

  """
  tape_data = collections.OrderedDict({})

  def record(f, *args, **kwargs):
    """Records execution to a tape."""
    name = kwargs.get("name")
    output = interceptable(f)(*args, **kwargs)
    if name:
      tape_data[name] = output
    return output

  with interception(record):
    yield tape_data

def
tape
(
)
:
tape_data
=
collections
.
OrderedDict
(
{
}
)
def
record
(
f
,
*
args
,
*
*
kwargs
)
:
"""Records execution to a tape."""
name
=
kwargs
.
get
(
"name"
)
output
=
interceptable
(
f
)
(
*
args
,
*
*
kwargs
)
if
name
:
tape_data
[
name
]
=
output
return
output
with
interception
(
record
)
:
yield
tape_data

Context manager for recording interceptable executions onto a tape .


def toy_logistic_data(num_examples, input_size=2, weights_prior_stddev=5.0):
  """Generates synthetic data for binary classification.

  Args:
    num_examples: The number of samples to generate (scalar Python `int`).
    input_size: The input space dimension (scalar Python `int`).
    weights_prior_stddev: The prior standard deviation of the weight
      vector. (scalar Python `float`).

  Returns:
    random_weights: Sampled weights as a Numpy `array` of shape
      `[input_size]`.
    random_bias: Sampled bias as a scalar Python `float`.
    design_matrix: Points sampled uniformly from the cube `[-1,
       1]^{input_size}`, as a Numpy `array` of shape `(num_examples,
       input_size)`.
    labels: Labels sampled from the logistic model `p(label=1) =
      logistic(dot(features, random_weights) + random_bias)`, as a Numpy
      `int32` `array` of shape `(num_examples, 1)`.
  """
  random_weights = weights_prior_stddev * np.random.randn(input_size)
  random_bias = np.random.randn()
  design_matrix = np.random.rand(num_examples, input_size) * 2 - 1
  logits = np.reshape(
      np.dot(design_matrix, random_weights) + random_bias,
      (-1, 1))
  p_labels = 1. / (1 + np.exp(-logits))
  labels = np.int32(p_labels > np.random.rand(num_examples, 1))
  return random_weights, random_bias, np.float32(design_matrix), labels

def toy_logistic_data ( num_examples , input_size = 2 , weights_prior_stddev = 5.0 ) : random_weights = weights_prior_stddev * np . random . randn ( input_size ) random_bias = np . random . randn ( ) design_matrix = np . random . rand ( num_examples , input_size ) * 2 - 1 logits = np . reshape ( np . dot ( design_matrix , random_weights ) + random_bias , ( - 1 , 1 ) ) p_labels = 1. / ( 1 + np . exp ( - logits ) ) labels = np . int32 ( p_labels > np . random . rand ( num_examples , 1 ) ) return random_weights , random_bias , np . float32 ( design_matrix ) , labels

Generates synthetic data for binary classification .


def visualize_decision(features, labels, true_w_b, candidate_w_bs, fname):
  """Utility method to visualize decision boundaries in R^2.

  Args:
    features: Input points, as a Numpy `array` of shape `[num_examples, 2]`.
    labels: Numpy `float`-like array of shape `[num_examples, 1]` giving a
      label for each point.
    true_w_b: A `tuple` `(w, b)` where `w` is a Numpy array of
       shape `[2]` and `b` is a scalar `float`, interpreted as a
       decision rule of the form `dot(features, w) + b > 0`.
    candidate_w_bs: Python `iterable` containing tuples of the same form as
       true_w_b.
    fname: The filename to save the plot as a PNG image (Python `str`).
  """
  fig = figure.Figure(figsize=(6, 6))
  canvas = backend_agg.FigureCanvasAgg(fig)
  ax = fig.add_subplot(1, 1, 1)
  ax.scatter(features[:, 0], features[:, 1],
             c=np.float32(labels[:, 0]),
             cmap=cm.get_cmap("binary"),
             edgecolors="k")

  def plot_weights(w, b, **kwargs):
    w1, w2 = w
    x1s = np.linspace(-1, 1, 100)
    x2s = -(w1  * x1s + b) / w2
    ax.plot(x1s, x2s, **kwargs)

  for w, b in candidate_w_bs:
    plot_weights(w, b,
                 alpha=1./np.sqrt(len(candidate_w_bs)),
                 lw=1, color="blue")

  if true_w_b is not None:
    plot_weights(*true_w_b, lw=4,
                 color="green", label="true separator")

  ax.set_xlim([-1.5, 1.5])
  ax.set_ylim([-1.5, 1.5])
  ax.legend()

  canvas.print_figure(fname, format="png")
  print("saved {}".format(fname))

def
visualize_decision
(
features
,
labels
,
true_w_b
,
candidate_w_bs
,
fname
)
:
fig
=
figure
.
Figure
(
figsize
=
(
6
,
6
)
)
canvas
=
backend_agg
.
FigureCanvasAgg
(
fig
)
ax
=
fig
.
add_subplot
(
1
,
1
,
1
)
ax
.
scatter
(
features
[
:
,
0
]
,
features
[
:
,
1
]
,
c
=
np
.
float32
(
labels
[
:
,
0
]
)
,
cmap
=
cm
.
get_cmap
(
"binary"
)
,
edgecolors
=
"k"
)
def
plot_weights
(
w
,
b
,
*
*
kwargs
)
:
w1
,
w2
=
w
x1s
=
np
.
linspace
(
1
,
1
,
100
)
x2s
=
(
w1
*
x1s
+
b
)
/
w2
ax
.
plot
(
x1s
,
x2s
,
*
*
kwargs
)
for
w
,
b
in
candidate_w_bs
:
plot_weights
(
w
,
b
,
alpha
=
1.
/
np
.
sqrt
(
len
(
candidate_w_bs
)
)
,
lw
=
1
,
color
=
"blue"
)
if
true_w_b
is
not
:
plot_weights
(
*
true_w_b
,
lw
=
4
,
color
=
"green"
,
label
=
"true separator"
)
ax
.
set_xlim
(
[
1.5
,
1.5
]
)
ax
.
set_ylim
(
[
1.5
,
1.5
]
)
ax
.
legend
(
)
canvas
.
print_figure
(
fname
,
format
=
"png"
)
print
(
"saved {}"
.
format
(
fname
)
)

Utility method to visualize decision boundaries in R^2 .


def build_input_pipeline(x, y, batch_size):
  """Build a Dataset iterator for supervised classification.

  Args:
    x: Numpy `array` of features, indexed by the first dimension.
    y: Numpy `array` of labels, with the same first dimension as `x`.
    batch_size: Number of elements in each training batch.

  Returns:
    batch_features: `Tensor` feed  features, of shape
      `[batch_size] + x.shape[1:]`.
    batch_labels: `Tensor` feed of labels, of shape
      `[batch_size] + y.shape[1:]`.
  """
  training_dataset = tf.data.Dataset.from_tensor_slices((x, y))
  training_batches = training_dataset.repeat().batch(batch_size)
  training_iterator = tf.compat.v1.data.make_one_shot_iterator(training_batches)
  batch_features, batch_labels = training_iterator.get_next()
  return batch_features, batch_labels

def build_input_pipeline ( x , y , batch_size ) : training_dataset = tf . data . Dataset . from_tensor_slices ( ( x , y ) ) training_batches = training_dataset . repeat ( ) . batch ( batch_size ) training_iterator = tf . compat . v1 . data . make_one_shot_iterator ( training_batches ) batch_features , batch_labels = training_iterator . get_next ( ) return batch_features , batch_labels

Build a Dataset iterator for supervised classification .


def _maybe_check_valid_map_values(map_values, validate_args):
  """Validate `map_values` if `validate_args`==True."""
  assertions = []

  message = 'Rank of map_values must be 1.'
  if tensorshape_util.rank(map_values.shape) is not None:
    if tensorshape_util.rank(map_values.shape) != 1:
      raise ValueError(message)
  elif validate_args:
    assertions.append(assert_util.assert_rank(map_values, 1, message=message))

  message = 'Size of map_values must be greater than 0.'
  if tensorshape_util.num_elements(map_values.shape) is not None:
    if tensorshape_util.num_elements(map_values.shape) == 0:
      raise ValueError(message)
  elif validate_args:
    assertions.append(
        assert_util.assert_greater(
            tf.size(input=map_values), 0, message=message))

  if validate_args:
    assertions.append(
        assert_util.assert_equal(
            tf.math.is_strictly_increasing(map_values),
            True,
            message='map_values is not strictly increasing.'))

  return assertions

Validate map_values if validate_args == True .


def trace(state: State, fn: TransitionOperator, num_steps: IntTensor,
          trace_fn: Callable[[State, TensorNest], TensorNest]
         ) -> Tuple[State, TensorNest]:
  """`TransitionOperator` that runs `fn` repeatedly and traces its outputs.

  Args:
    state: A nest of `Tensor`s or None.
    fn: A `TransitionOperator`.
    num_steps: Number of steps to run the function for. Must be greater than 1.
    trace_fn: Callable that the unpacked outputs of `fn` and returns a nest of
      `Tensor`s. These will be stacked and returned.

  Returns:
    state: The final state returned by `fn`.
    traces: Stacked outputs of `trace_fn`.
  """

  def fn_wrapper(args, _):
    return tf.nest.map_structure(tf.convert_to_tensor, call_fn(fn, args[0]))

  def trace_fn_wrapper(args):
    return tf.nest.map_structure(tf.convert_to_tensor, call_fn(trace_fn, args))

  state = call_fn(fn, state)
  first_trace = trace_fn_wrapper(state)

  state, full_trace = mcmc_util.trace_scan(
      fn_wrapper, state, tf.ones(num_steps - 1), trace_fn=trace_fn_wrapper)

  prepend = lambda x, y: tf.concat(  # pylint: disable=g-long-lambda
      [tf.convert_to_tensor(value=x)[tf.newaxis], y], 0)

  return state, tf.nest.map_structure(prepend, first_trace, full_trace)

def
trace
(
state
:
State
,
fn
:
TransitionOperator
,
num_steps
:
IntTensor
,
trace_fn
:
Callable
[
[
State
,
TensorNest
]
,
TensorNest
]
)
->
Tuple
[
State
,
TensorNest
]
:
def
fn_wrapper
(
args
,
_
)
:
return
tf
.
nest
.
map_structure
(
tf
.
convert_to_tensor
,
call_fn
(
fn
,
args
[
0
]
)
)
def
trace_fn_wrapper
(
args
)
:
return
tf
.
nest
.
map_structure
(
tf
.
convert_to_tensor
,
call_fn
(
trace_fn
,
args
)
)
state
=
call_fn
(
fn
,
state
)
first_trace
=
trace_fn_wrapper
(
state
)
state
,
full_trace
=
mcmc_util
.
trace_scan
(
fn_wrapper
,
state
,
tf
.
ones
(
num_steps
1
)
,
trace_fn
=
trace_fn_wrapper
)
prepend
=
lambda
x
,
y
:
tf
.
concat
(
# pylint: disable=g-long-lambda
[
tf
.
convert_to_tensor
(
value
=
x
)
[
tf
.
newaxis
]
,
y
]
,
0
)
return
state
,
tf
.
nest
.
map_structure
(
prepend
,
first_trace
,
full_trace
)

TransitionOperator that runs fn repeatedly and traces its outputs .


def call_fn(fn: TransitionOperator, args: Union[Tuple[Any], Any]) -> Any:
  """Calls a transition operator with args, unpacking args if its a sequence.

  Args:
    fn: A `TransitionOperator`.
    args: Arguments to `fn`

  Returns:
    ret: Return value of `fn`.
  """

  if isinstance(args, (list, tuple)) and not mcmc_util.is_namedtuple_like(args):
    args = args  # type: Tuple[Any]
    return fn(*args)
  else:
    return fn(args)

def
call_fn
(
fn
:
TransitionOperator
,
args
:
Union
[
Tuple
[
Any
]
,
Any
]
)
->
Any
:
if
isinstance
(
args
,
(
list
,
tuple
)
)
and
not
mcmc_util
.
is_namedtuple_like
(
args
)
:
args
=
args
# type: Tuple[Any]
return
fn
(
*
args
)
else
:
return
fn
(
args
)

Calls a transition operator with args unpacking args if its a sequence .


def call_and_grads(fn: TransitionOperator, args: Union[Tuple[Any], Any]
                  ) -> Tuple[tf.Tensor, TensorNest, TensorNest]:
  """Calls `fn` and returns the gradients with respect to `fn`'s first output.

  Args:
    fn: A `TransitionOperator`.
    args: Arguments to `fn`

  Returns:
    ret: First output of `fn`.
    extra: Second output of `fn`.
    grads: Gradients of `ret` with respect to `args`.
  """
  with tf.GradientTape() as tape:
    tape.watch(args)
    ret, extra = call_fn(fn, args)
  grads = tape.gradient(ret, args)
  return ret, extra, grads

def call_and_grads ( fn : TransitionOperator , args : Union [ Tuple [ Any ] , Any ] ) -> Tuple [ tf . Tensor , TensorNest , TensorNest ] : with tf . GradientTape ( ) as tape : tape . watch ( args ) ret , extra = call_fn ( fn , args ) grads = tape . gradient ( ret , args ) return ret , extra , grads

Calls fn and returns the gradients with respect to fn s first output .


def maybe_broadcast_structure(from_structure: Any, to_structure: Any) -> Any:
  """Maybe broadcasts `from_structure` to `to_structure`.

  If `from_structure` is a singleton, it is tiled to match the structure of
  `to_structure`. Note that the elements in `from_structure` are not copied if
  this tiling occurs.

  Args:
    from_structure: A structure.
    to_structure: A structure.

  Returns:
    new_from_structure: Same structure as `to_structure`.
  """
  flat_from = tf.nest.flatten(from_structure)
  flat_to = tf.nest.flatten(to_structure)
  if len(flat_from) == 1:
    flat_from *= len(flat_to)
  return tf.nest.pack_sequence_as(to_structure, flat_from)

def maybe_broadcast_structure ( from_structure : Any , to_structure : Any ) -> Any : flat_from = tf . nest . flatten ( from_structure ) flat_to = tf . nest . flatten ( to_structure ) if len ( flat_from ) == 1 : flat_from *= len ( flat_to ) return tf . nest . pack_sequence_as ( to_structure , flat_from )

Maybe broadcasts from_structure to to_structure .


def transform_log_prob_fn(log_prob_fn: PotentialFn,
                          bijector: BijectorNest,
                          init_state: State = None
                         ) -> Union[PotentialFn, Tuple[PotentialFn, State]]:
  """Transforms a log-prob function using a bijector.

  This takes a log-prob function and creates a new log-prob function that now
  takes takes state in the domain of the bijector, forward transforms that state
  and calls the original log-prob function. It then returns the log-probability
  that correctly accounts for this transformation.

  The forward-transformed state is pre-pended to the original log-prob
  function's extra returns and returned as the new extra return.

  For convenience you can also pass the initial state (in the original space),
  and this function will return the inverse transformed as the 2nd return value.
  You'd use this to initialize MCMC operators that operate in the transformed
  space.

  Args:
    log_prob_fn: Log prob fn.
    bijector: Bijector(s), must be of the same structure as the `log_prob_fn`
      inputs.
    init_state: Initial state, in the original space.

  Returns:
    transformed_log_prob_fn: Transformed log prob fn.
    transformed_init_state: If `init_state` is provided. Initial state in the
      transformed space.
  """

  def wrapper(*args):
    """Transformed wrapper."""
    bijector_ = bijector

    args = tf.nest.map_structure(lambda x: 0. + x, args)
    if len(args) == 1:
      args = args[0]
    elif isinstance(bijector_, list):
      bijector_ = tuple(bijector_)

    original_space_args = tf.nest.map_structure(lambda b, x: b.forward(x),
                                                bijector_, args)
    original_space_args = original_space_args  # type: Tuple[Any]
    original_space_log_prob, extra = call_fn(log_prob_fn, original_space_args)
    event_ndims = tf.nest.map_structure(
        lambda x: tf.rank(x) - tf.rank(original_space_log_prob), args)

    return original_space_log_prob + sum(
        tf.nest.flatten(
            tf.nest.map_structure(
                lambda b, x, e: b.forward_log_det_jacobian(x, event_ndims=e),
                bijector_, args, event_ndims))), [original_space_args, extra]

  if init_state is None:
    return wrapper
  else:
    return wrapper, tf.nest.map_structure(lambda b, s: b.inverse(s), bijector,
                                          init_state)

def
transform_log_prob_fn
(
log_prob_fn
:
PotentialFn
,
bijector
:
BijectorNest
,
init_state
:
State
=
)
->
Union
[
PotentialFn
,
Tuple
[
PotentialFn
,
State
]
]
:
def
wrapper
(
*
args
)
:
"""Transformed wrapper."""
bijector_
=
bijector
args
=
tf
.
nest
.
map_structure
(
lambda
x
:
0.
+
x
,
args
)
if
len
(
args
)
==
1
:
args
=
args
[
0
]
elif
isinstance
(
bijector_
,
list
)
:
bijector_
=
tuple
(
bijector_
)
original_space_args
=
tf
.
nest
.
map_structure
(
lambda
b
,
x
:
b
.
forward
(
x
)
,
bijector_
,
args
)
original_space_args
=
original_space_args
# type: Tuple[Any]
original_space_log_prob
,
extra
=
call_fn
(
log_prob_fn
,
original_space_args
)
event_ndims
=
tf
.
nest
.
map_structure
(
lambda
x
:
tf
.
rank
(
x
)
tf
.
rank
(
original_space_log_prob
)
,
args
)
return
original_space_log_prob
+
sum
(
tf
.
nest
.
flatten
(
tf
.
nest
.
map_structure
(
lambda
b
,
x
,
e
:
b
.
forward_log_det_jacobian
(
x
,
event_ndims
=
e
)
,
bijector_
,
args
,
event_ndims
)
)
)
,
[
original_space_args
,
extra
]
if
init_state
is
:
return
wrapper
else
:
return
wrapper
,
tf
.
nest
.
map_structure
(
lambda
b
,
s
:
b
.
inverse
(
s
)
,
bijector
,
init_state
)

Transforms a log - prob function using a bijector .


def leapfrog_step(leapfrog_step_state: LeapFrogStepState,
                  step_size: FloatTensor, target_log_prob_fn: PotentialFn,
                  kinetic_energy_fn: PotentialFn
                 ) -> Tuple[LeapFrogStepState, LeapFrogStepExtras]:
  """Leapfrog `TransitionOperator`.

  Args:
    leapfrog_step_state: LeapFrogStepState.
    step_size: Step size, structure broadcastable to the `target_log_prob_fn`
      state.
    target_log_prob_fn: Target log prob fn.
    kinetic_energy_fn: Kinetic energy fn.

  Returns:
    leapfrog_step_state: LeapFrogStepState.
    leapfrog_step_extras: LeapFrogStepExtras.
  """
  state = leapfrog_step_state.state
  state_grads = leapfrog_step_state.state_grads
  momentum = leapfrog_step_state.momentum
  step_size = maybe_broadcast_structure(step_size, state)

  state = tf.nest.map_structure(tf.convert_to_tensor, state)
  momentum = tf.nest.map_structure(tf.convert_to_tensor, momentum)
  state = tf.nest.map_structure(tf.convert_to_tensor, state)

  if state_grads is None:
    _, _, state_grads = call_and_grads(target_log_prob_fn, state)
  else:
    state_grads = tf.nest.map_structure(tf.convert_to_tensor, state_grads)

  momentum = tf.nest.map_structure(lambda m, sg, s: m + 0.5 * sg * s, momentum,
                                   state_grads, step_size)

  kinetic_energy, kinetic_energy_extra, momentum_grads = call_and_grads(
      kinetic_energy_fn, momentum)

  state = tf.nest.map_structure(lambda x, mg, s: x + mg * s, state,
                                momentum_grads, step_size)

  target_log_prob, state_extra, state_grads = call_and_grads(
      target_log_prob_fn, state)

  momentum = tf.nest.map_structure(lambda m, sg, s: m + 0.5 * sg * s, momentum,
                                   state_grads, step_size)

  return LeapFrogStepState(state, state_grads, momentum), LeapFrogStepExtras(
      target_log_prob, state_extra, kinetic_energy, kinetic_energy_extra)

def leapfrog_step ( leapfrog_step_state : LeapFrogStepState , step_size : FloatTensor , target_log_prob_fn : PotentialFn , kinetic_energy_fn : PotentialFn ) -> Tuple [ LeapFrogStepState , LeapFrogStepExtras ] : state = leapfrog_step_state . state state_grads = leapfrog_step_state . state_grads momentum = leapfrog_step_state . momentum step_size = maybe_broadcast_structure ( step_size , state ) state = tf . nest . map_structure ( tf . convert_to_tensor , state ) momentum = tf . nest . map_structure ( tf . convert_to_tensor , momentum ) state = tf . nest . map_structure ( tf . convert_to_tensor , state ) if state_grads is None : _ , _ , state_grads = call_and_grads ( target_log_prob_fn , state ) else : state_grads = tf . nest . map_structure ( tf . convert_to_tensor , state_grads ) momentum = tf . nest . map_structure ( lambda m , sg , s : m + 0.5 * sg * s , momentum , state_grads , step_size ) kinetic_energy , kinetic_energy_extra , momentum_grads = call_and_grads ( kinetic_energy_fn , momentum ) state = tf . nest . map_structure ( lambda x , mg , s : x + mg * s , state , momentum_grads , step_size ) target_log_prob , state_extra , state_grads = call_and_grads ( target_log_prob_fn , state ) momentum = tf . nest . map_structure ( lambda m , sg , s : m + 0.5 * sg * s , momentum , state_grads , step_size ) return LeapFrogStepState ( state , state_grads , momentum ) , LeapFrogStepExtras ( target_log_prob , state_extra , kinetic_energy , kinetic_energy_extra )


def metropolis_hastings_step(current_state: State,
                             proposed_state: State,
                             energy_change: FloatTensor,
                             seed=None) -> Tuple[State, tf.Tensor, tf.Tensor]:
  """Metropolis-Hastings step.

  This probabilistically chooses between `current_state` and `proposed_state`
  based on the `energy_change` so as to preserve detailed balance.

  Energy change is the negative of `log_accept_ratio`.

  Args:
    current_state: Current state.
    proposed_state: Proposed state.
    energy_change: E(proposed_state) - E(previous_state).
    seed: For reproducibility.

  Returns:
    new_state: The chosen state.
    is_accepted: Whether the proposed state was accepted.
    log_uniform: The random number that was used to select between the two
      states.
  """
  flat_current = tf.nest.flatten(current_state)
  flat_proposed = nest.flatten_up_to(current_state, proposed_state)
  # Impute the None's in the current state.
  flat_current = [
      p if c is None else c for p, c in zip(flat_proposed, flat_current)
  ]
  current_state = tf.nest.pack_sequence_as(current_state, flat_current)

  current_state = tf.nest.map_structure(tf.convert_to_tensor, current_state)
  proposed_state = tf.nest.map_structure(tf.convert_to_tensor, proposed_state)
  energy_change = tf.convert_to_tensor(value=energy_change)

  log_accept_ratio = -energy_change

  log_uniform = tf.math.log(
      tf.random.uniform(
          shape=tf.shape(input=log_accept_ratio),
          dtype=log_accept_ratio.dtype.base_dtype,
          seed=seed))
  is_accepted = log_uniform < log_accept_ratio

  next_state = mcmc_util.choose(
      is_accepted, proposed_state, current_state, name='choose_next_state')
  return next_state, is_accepted, log_uniform


def hamiltonian_monte_carlo(
    hmc_state: HamiltonianMonteCarloState,
    target_log_prob_fn: PotentialFn,
    step_size: Any,
    num_leapfrog_steps: IntTensor,
    momentum: State = None,
    kinetic_energy_fn: PotentialFn = None,
    momentum_sample_fn: MomentumSampleFn = None,
    leapfrog_trace_fn: Callable[[LeapFrogStepState, LeapFrogStepExtras],
                                TensorNest] = lambda *args: (),
    seed=None,
) -> Tuple[HamiltonianMonteCarloState, HamiltonianMonteCarloExtra]:
  """Hamiltonian Monte Carlo `TransitionOperator`.

  #### Example

  ```python
  step_size = 0.2
  num_steps = 2000
  num_leapfrog_steps = 10
  state = tf.ones([16, 2])

  base_mean = [1., 0]
  base_cov = [[1, 0.5], [0.5, 1]]

  bijector = tfb.Softplus()
  base_dist = tfd.MultivariateNormalFullCovariance(
      loc=base_mean, covariance_matrix=base_cov)
  target_dist = bijector(base_dist)

  def orig_target_log_prob_fn(x):
    return target_dist.log_prob(x), ()

  target_log_prob_fn, state = fun_mcmc.transform_log_prob_fn(
      orig_target_log_prob_fn, bijector, state)

  kernel = tf.function(lambda state: fun_mcmc.hamiltonian_monte_carlo(
      state,
      step_size=step_size,
      num_leapfrog_steps=num_leapfrog_steps,
      target_log_prob_fn=target_log_prob_fn,
      seed=tfp_test_util.test_seed()))

  _, chain = fun_mcmc.trace(
      state=fun_mcmc.HamiltonianMonteCarloState(
          state=state,
          state_grads=None,
          target_log_prob=None,
          state_extra=None),
      fn=kernel,
      num_steps=num_steps,
      trace_fn=lambda state, extra: state.state_extra[0])
  ```

  Args:
    hmc_state: HamiltonianMonteCarloState.
    target_log_prob_fn: Target log prob fn.
    step_size: Step size, structure broadcastable to the `target_log_prob_fn`
      state.
    num_leapfrog_steps: Number of leapfrog steps to take.
    momentum: Initial momentum, passed to `momentum_sample_fn`. Default: zeroes.
    kinetic_energy_fn: Kinetic energy function.
    momentum_sample_fn: Sampler for the momentum.
    leapfrog_trace_fn: Trace function for the leapfrog integrator.
    seed: For reproducibility.

  Returns:
    hmc_state: HamiltonianMonteCarloState
    hmc_extra: HamiltonianMonteCarloExtra
  """
  state = hmc_state.state
  state_grads = hmc_state.state_grads
  target_log_prob = hmc_state.target_log_prob
  state_extra = hmc_state.state_extra

  if kinetic_energy_fn is None:

    # pylint: disable=function-redefined
    def kinetic_energy_fn(*momentum):
      return tf.add_n([
          tf.reduce_sum(input_tensor=tf.square(x), axis=-1) / 2.
          for x in tf.nest.flatten(momentum)
      ]), ()

  if momentum_sample_fn is None:

    # pylint: disable=function-redefined
    def momentum_sample_fn(*momentum):
      ret = tf.nest.map_structure(
          lambda x: tf.random.normal(tf.shape(input=x), dtype=x.dtype),
          momentum)
      if len(ret) == 1:
        return ret[0]
      else:
        return ret

  if momentum is None:
    momentum = call_fn(momentum_sample_fn,
                       tf.nest.map_structure(tf.zeros_like, state))
  if target_log_prob is None:
    target_log_prob, state_extra, state_grads = call_and_grads(
        target_log_prob_fn, state)

  kinetic_energy, _ = call_fn(kinetic_energy_fn, momentum)
  current_energy = -target_log_prob + kinetic_energy
  current_state = HamiltonianMonteCarloState(
      state=state,
      state_grads=state_grads,
      state_extra=state_extra,
      target_log_prob=target_log_prob)

  def leapfrog_wrapper(leapfrog_state, target_log_prob, state_extra):
    """Leapfrog wrapper that tracks extra state."""
    del target_log_prob
    del state_extra

    leapfrog_state, leapfrog_extra = leapfrog_step(
        leapfrog_state,
        step_size=step_size,
        target_log_prob_fn=target_log_prob_fn,
        kinetic_energy_fn=kinetic_energy_fn)

    return [
        leapfrog_state, leapfrog_extra.target_log_prob,
        leapfrog_extra.state_extra
    ], leapfrog_extra

  def leapfrog_trace_wrapper_fn(args, leapfrog_extra):
    return leapfrog_trace_fn(args[0], leapfrog_extra)

  leapfrog_wrapper_state = (LeapFrogStepState(state, state_grads, momentum),
                            target_log_prob, state_extra)

  [[leapfrog_state, target_log_prob, state_extra], _], leapfrog_trace = trace(
      leapfrog_wrapper_state,
      leapfrog_wrapper,
      num_leapfrog_steps,
      trace_fn=leapfrog_trace_wrapper_fn)

  kinetic_energy, _ = call_fn(kinetic_energy_fn, leapfrog_state.momentum)
  proposed_energy = -target_log_prob + kinetic_energy
  proposed_state = HamiltonianMonteCarloState(
      state=leapfrog_state.state,
      state_grads=leapfrog_state.state_grads,
      target_log_prob=target_log_prob,
      state_extra=state_extra)

  energy_change = proposed_energy - current_energy
  hmc_state, is_accepted, _ = metropolis_hastings_step(
      current_state, proposed_state, energy_change, seed=seed)

  hmc_state = hmc_state  # type: HamiltonianMonteCarloState
  return hmc_state, HamiltonianMonteCarloExtra(
      is_accepted=is_accepted,
      proposed_hmc_state=proposed_state,
      log_accept_ratio=-energy_change,
      leapfrog_trace=leapfrog_trace)

Hamiltonian Monte Carlo TransitionOperator .


def sign_adaptation(control: FloatNest,
                    output: FloatTensor,
                    set_point: FloatTensor,
                    adaptation_rate: FloatTensor = 0.01) -> FloatNest:
  """A function to do simple sign-based control of a variable.

  ```
  control = control * (1. + adaptation_rate) ** sign(output - set_point)
  ```

  Args:
    control: The control variable.
    output: The output variable.
    set_point: The set point for `output`. This function will adjust `control`
      so that `output` matches `set_point`.
    adaptation_rate: Adaptation rate.

  Returns:
    control: New control.
  """

  def _get_new_control(control, output, set_point):
    new_control = mcmc_util.choose(output > set_point,
                                   control * (1. + adaptation_rate),
                                   control / (1. + adaptation_rate))
    return new_control

  output = maybe_broadcast_structure(output, control)
  set_point = maybe_broadcast_structure(set_point, control)

  return tf.nest.map_structure(_get_new_control, control, output, set_point)

def sign_adaptation ( control : FloatNest , output : FloatTensor , set_point : FloatTensor , adaptation_rate : FloatTensor = 0.01 ) -> FloatNest : def _get_new_control ( control , output , set_point ) : new_control = mcmc_util . choose ( output > set_point , control * ( 1. + adaptation_rate ) , control / ( 1. + adaptation_rate ) ) return new_control output = maybe_broadcast_structure ( output , control ) set_point = maybe_broadcast_structure ( set_point , control ) return tf . nest . map_structure ( _get_new_control , control , output , set_point )

A function to do simple sign - based control of a variable .


def from_config(cls, config):
    """Creates a layer from its config.

    This method is the reverse of `get_config`, capable of instantiating the
    same layer from the config dictionary.

    Args:
      config: A Python dictionary, typically the output of `get_config`.

    Returns:
      layer: A layer instance.
    """
    config = config.copy()
    function_keys = [
        'kernel_posterior_fn',
        'kernel_posterior_tensor_fn',
        'kernel_prior_fn',
        'kernel_divergence_fn',
        'bias_posterior_fn',
        'bias_posterior_tensor_fn',
        'bias_prior_fn',
        'bias_divergence_fn',
    ]
    for function_key in function_keys:
      serial = config[function_key]
      function_type = config.pop(function_key + '_type')
      if serial is not None:
        config[function_key] = tfp_layers_util.deserialize_function(
            serial,
            function_type=function_type)
    return cls(**config)

def from_config ( cls , config ) : config = config . copy ( ) function_keys = [ 'kernel_posterior_fn' , 'kernel_posterior_tensor_fn' , 'kernel_prior_fn' , 'kernel_divergence_fn' , 'bias_posterior_fn' , 'bias_posterior_tensor_fn' , 'bias_prior_fn' , 'bias_divergence_fn' , ] for function_key in function_keys : serial = config [ function_key ] function_type = config . pop ( function_key + '_type' ) if serial is not None : config [ function_key ] = tfp_layers_util . deserialize_function ( serial , function_type = function_type ) return cls ( * * config )


def _as_tensor(x, name, dtype):
  """Convenience to convert to `Tensor` or leave as `None`."""
  return None if x is None else tf.convert_to_tensor(
      value=x, name=name, dtype=dtype)

def _as_tensor ( x , name , dtype ) : return None if x is None else tf . convert_to_tensor ( value = x , name = name , dtype = dtype )

Convenience to convert to Tensor or leave as None .


def _create_scale_operator(self, identity_multiplier, diag, tril,
                             perturb_diag, perturb_factor, shift, validate_args,
                             dtype):
    """Construct `scale` from various components.

    Args:
      identity_multiplier: floating point rank 0 `Tensor` representing a scaling
        done to the identity matrix.
      diag: Floating-point `Tensor` representing the diagonal matrix.`diag` has
        shape `[N1, N2, ...  k]`, which represents a k x k diagonal matrix.
      tril: Floating-point `Tensor` representing the lower triangular matrix.
       `tril` has shape `[N1, N2, ...  k, k]`, which represents a k x k lower
       triangular matrix.
      perturb_diag: Floating-point `Tensor` representing the diagonal matrix of
        the low rank update.
      perturb_factor: Floating-point `Tensor` representing factor matrix.
      shift: Floating-point `Tensor` representing `shift in `scale @ X + shift`.
      validate_args: Python `bool` indicating whether arguments should be
        checked for correctness.
      dtype: `DType` for arg `Tensor` conversions.

    Returns:
      scale. In the case of scaling by a constant, scale is a
      floating point `Tensor`. Otherwise, scale is a `LinearOperator`.

    Raises:
      ValueError: if all of `tril`, `diag` and `identity_multiplier` are `None`.
    """
    identity_multiplier = _as_tensor(identity_multiplier, "identity_multiplier",
                                     dtype)
    diag = _as_tensor(diag, "diag", dtype)
    tril = _as_tensor(tril, "tril", dtype)
    perturb_diag = _as_tensor(perturb_diag, "perturb_diag", dtype)
    perturb_factor = _as_tensor(perturb_factor, "perturb_factor", dtype)

    # If possible, use the low rank update to infer the shape of
    # the identity matrix, when scale represents a scaled identity matrix
    # with a low rank update.
    shape_hint = None
    if perturb_factor is not None:
      shape_hint = distribution_util.dimension_size(perturb_factor, axis=-2)

    if self._is_only_identity_multiplier:
      if validate_args:
        return distribution_util.with_dependencies([
            assert_util.assert_none_equal(
                identity_multiplier, tf.zeros([], identity_multiplier.dtype),
                ["identity_multiplier should be non-zero."])
        ], identity_multiplier)
      return identity_multiplier

    scale = distribution_util.make_tril_scale(
        loc=shift,
        scale_tril=tril,
        scale_diag=diag,
        scale_identity_multiplier=identity_multiplier,
        validate_args=validate_args,
        assert_positive=False,
        shape_hint=shape_hint)

    if perturb_factor is not None:
      return tf.linalg.LinearOperatorLowRankUpdate(
          scale,
          u=perturb_factor,
          diag_update=perturb_diag,
          is_diag_update_positive=perturb_diag is None,
          is_non_singular=True,  # Implied by is_positive_definite=True.
          is_self_adjoint=True,
          is_positive_definite=True,
          is_square=True)

    return scale

Construct scale from various components .