aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/tensorboard/components/tf-categorizer/categorizer.ts
blob: e05078279e268fb8e3edd2a9e569749271d589d9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
/// <reference path="../../typings/tsd.d.ts" />

module Categorizer {
  /**
   * This module contains methods that allow sorting tags into "categories".
   * A category contains a name and a list of tags.
   * The sorting strategy is defined by a "CustomCategorization", which contains
   * "categoryDefinitions" which are regex rules used to construct a category.
   * E.g. the regex rule "xent" will create a category called "xent" that
   * contains values whose tags match the regex.
   *
   * After custom categories are evaluated, the tags are sorted by a hardcoded
   * fallback categorizer, which may, for example, group tags into categories
   * based on their top namespace.
   */

  export interface Category {
    // Categories that data is sorted into
    name: string;
    tags: string[];
  }

  export interface CustomCategorization {
    // Defines a categorization strategy
    categoryDefinitions: string[];
    fallbackCategorizer: string;
    /* {"TopLevelNamespaceCategorizer",
        "LegacyUnderscoreCategorizer"} */
  }

  export interface Categorizer {
    // Function that generates categories
    (tags: string[]): Category[];
  }

  /* Canonical TensorFlow ops are namespaced using forward slashes.
   * This fallback categorizer categorizes by the top-level namespace.
   */
  export var topLevelNamespaceCategorizer: Categorizer = splitCategorizer(/\//);

  // Try to produce good categorizations on legacy graphs, which often
  // are namespaced like l1_foo/bar or l2_baz/bam.
  // If there is no leading underscore before the first forward slash,
  // then it behaves the same as topLevelNamespaceCategorizer
  export var legacyUnderscoreCategorizer: Categorizer = splitCategorizer(/[\/_]/);

  export function fallbackCategorizer(s: string): Categorizer {
    switch (s) {
      case "TopLevelNamespaceCategorizer":
        return topLevelNamespaceCategorizer;
      case "LegacyUnderscoreCategorizer":
        return legacyUnderscoreCategorizer;
      default:
        throw new Error("Unrecognized categorization strategy: " + s);
    }
  }

  /* An "extractor" is a function that takes a tag name, and "extracts" a category name.
   * This function takes an extractor, and produces a categorizer.
   * Currently, it is just used for the fallbackCategorizer, but we may want to
   * refactor the general categorization logic to use the concept of extractors.
   */
  function extractorToCategorizer(extractor: (s: string) => string): Categorizer {
    return (tags: string[]): Category[] => {
      if (tags.length === 0) {
        return [];
      }
      var sortedTags = tags.slice().sort();
      var categories: Category[] = [];
      var currentCategory = {
        name: extractor(sortedTags[0]),
        tags: [],
      };
      sortedTags.forEach((t: string) => {
        var topLevel = extractor(t);
        if (currentCategory.name !== topLevel) {
          categories.push(currentCategory);
          currentCategory = {
            name: topLevel,
            tags: [],
          };
        }
        currentCategory.tags.push(t);
      });
      categories.push(currentCategory);
      return categories;
    };
  }

  function splitCategorizer(r: RegExp): Categorizer {
    var extractor = (t: string) => {
      return t.split(r)[0];
    };
    return extractorToCategorizer(extractor);
  }

  export interface CategoryDefinition {
    name: string;
    matches: (t: string) => boolean;
  }

  export function defineCategory(ruledef: string): CategoryDefinition {
    var r = new RegExp(ruledef);
    var f = function(tag: string): boolean {
      return r.test(tag);
    };
    return { name: ruledef, matches: f };
  }

  export function _categorizer(rules: CategoryDefinition[], fallback: Categorizer) {
    return function(tags: string[]): Category[] {
      var remaining: d3.Set = d3.set(tags);
      var userSpecified = rules.map((def: CategoryDefinition) => {
        var tags: string[] = [];
        remaining.forEach((t: string) => {
          if (def.matches(t)) {
            tags.push(t);
          }
        });
        var cat = { name: def.name, tags: tags.sort() };
        return cat;
      });
      var defaultCategories = fallback(remaining.values());
      return userSpecified.concat(defaultCategories);
    };
  }

  export function categorizer(s: CustomCategorization): Categorizer {
    var rules = s.categoryDefinitions.map(defineCategory);
    var fallback = fallbackCategorizer(s.fallbackCategorizer);
    return _categorizer(rules, fallback);
  };
}